作者|DR. VAIBHAV KUMAR
编译|VK
起源|Analytics In Diamag

随着深度学习模型在各种利用中的胜利施行，当初是时候取得不仅精确而且速度更快的后果。

为了失去更精确的后果，数据的大小是十分重要的，然而当这个大小影响到机器学习模型的训练工夫时，这始终是一个值得关注的问题。

为了克服训练工夫的问题，咱们应用TPU运行时环境来减速训练。为此，PyTorch始终在通过提供最先进的硬件加速器来反对机器学习的实现。

PyTorch对云TPU的反对是通过与XLA（减速线性代数）的集成实现的，XLA是一种用于线性代数的编译器，能够针对多种类型的硬件，包含CPU、GPU和TPU。

本文演示了如何应用PyTorch和TPU实现深度学习模型，以放慢训练过程。

在这里，咱们应用PyTorch定义了一个卷积神经网络（CNN）模型，并在PyTorch/XLA环境中对该模型进行了训练。

XLA将CNN模型与分布式多解决环境中的Google Cloud TPU（张量处理单元）连接起来。在这个实现中，应用8个TPU外围来创立一个多解决环境。

咱们将用这个PyTorch深度学习框架进行时装分类测试，察看训练工夫和准确性。

用PyTorch和TPU实现CNN

咱们将在Google Colab中实现执行，因为它提供收费的云TPU（张量处理单元）。

在持续下一步之前，在Colab笔记本中，转到“编辑”，而后抉择“设置”，从上面屏幕截图中的列表中抉择“TPU”作为“硬件加速器”。

验证TPU上面的代码是否失常运行。

import osassert os.environ['COLAB_TPU_ADDR']

如果启用了TPU，它将胜利执行，否则它将返回‘KeyError: ‘COLAB_TPU_ADDR’’。你也能够通过打印TPU地址来查看TPU。

TPU_Path = 'grpc://'+os.environ['COLAB_TPU_ADDR']print('TPU Address:', TPU_Path)

启用TPU后，咱们将装置兼容的管制盘和依赖项，以应用以下代码设置XLA环境。

VERSION = "20200516" !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py!python pytorch-xla-env-setup.py --version $VERSION

一旦装置胜利，咱们将持续定义加载数据集、初始化CNN模型、训练和测试的办法。首先，咱们将导入所需的库。

import numpy as npimport osimport timeimport torchimport torch.nn as nnimport torch.nn.functional as Fimport torch.optim as optimimport torch_xlaimport torch_xla.core.xla_model as xmimport torch_xla.debug.metrics as metimport torch_xla.distributed.parallel_loader as plimport torch_xla.distributed.xla_multiprocessing as xmpimport torch_xla.utils.utils as xufrom torchvision import datasets, transforms

之后，咱们将进一步定义须要的超参数。

# 定义参数FLAGS = {}FLAGS['datadir'] = "/tmp/mnist"FLAGS['batch_size'] = 128FLAGS['num_workers'] = 4FLAGS['learning_rate'] = 0.01FLAGS['momentum'] = 0.5FLAGS['num_epochs'] = 50FLAGS['num_cores'] = 8FLAGS['log_steps'] = 20FLAGS['metrics_debug'] = False

上面的代码片段将把CNN模型定义为PyTorch实例，以及用于加载数据、训练模型和测试模型的函数。

SERIAL_EXEC = xmp.MpSerialExecutor()class FashionMNIST(nn.Module):  def __init__(self):    super(FashionMNIST, self).__init__()    self.conv1 = nn.Conv2d(1, 10, kernel_size=5)    self.bn1 = nn.BatchNorm2d(10)    self.conv2 = nn.Conv2d(10, 20, kernel_size=5)    self.bn2 = nn.BatchNorm2d(20)    self.fc1 = nn.Linear(320, 50)    self.fc2 = nn.Linear(50, 10)  def forward(self, x):    x = F.relu(F.max_pool2d(self.conv1(x), 2))    x = self.bn1(x)    x = F.relu(F.max_pool2d(self.conv2(x), 2))    x = self.bn2(x)    x = torch.flatten(x, 1)    x = F.relu(self.fc1(x))    x = self.fc2(x)    return F.log_softmax(x, dim=1)# 只在内存中实例化一次模型权重。WRAPPED_MODEL = xmp.MpModelWrapper(FashionMNIST())def train_mnist():  torch.manual_seed(1)   def get_dataset():    norm = transforms.Normalize((0.1307,), (0.3081,))    train_dataset = datasets.FashionMNIST(        FLAGS['datadir'],        train=True,        download=True,        transform=transforms.Compose(            [transforms.ToTensor(), norm]))    test_dataset = datasets.FashionMNIST(        FLAGS['datadir'],        train=False,        download=True,        transform=transforms.Compose(            [transforms.ToTensor(), norm]))      return train_dataset, test_dataset  #应用串行执行器能够防止多个过程下载雷同的数据  train_dataset, test_dataset = SERIAL_EXEC.run(get_dataset)  train_sampler = torch.utils.data.distributed.DistributedSampler(    train_dataset,    num_replicas=xm.xrt_world_size(),    rank=xm.get_ordinal(),    shuffle=True)  train_loader = torch.utils.data.DataLoader(      train_dataset,      batch_size=FLAGS['batch_size'],      sampler=train_sampler,      num_workers=FLAGS['num_workers'],      drop_last=True)  test_loader = torch.utils.data.DataLoader(      test_dataset,      batch_size=FLAGS['batch_size'],      shuffle=False,      num_workers=FLAGS['num_workers'],      drop_last=True)  # 调整学习率  lr = FLAGS['learning_rate'] * xm.xrt_world_size()  # 获取损失函数、优化器和模型  device = xm.xla_device()  model = WRAPPED_MODEL.to(device)  optimizer = optim.SGD(model.parameters(), lr=lr, momentum=FLAGS['momentum'])  loss_fn = nn.NLLLoss()  def train_fun(loader):    tracker = xm.RateTracker()    model.train()    for x, (data, target) in enumerate(loader):      optimizer.zero_grad()      output = model(data)      loss = loss_fn(output, target)      loss.backward()      xm.optimizer_step(optimizer)      tracker.add(FLAGS['batch_size'])      if x % FLAGS['log_steps'] == 0:        print('[xla:{}]({}) Loss={:.5f}'.format(            xm.get_ordinal(), x, loss.item(), time.asctime()), flush=True)  def test_fun(loader):    total_samples = 0    correct = 0    model.eval()    data, pred, target = None, None, None    for data, target in loader:      output = model(data)      pred = output.max(1, keepdim=True)[1]      correct += pred.eq(target.view_as(pred)).sum().item()      total_samples += data.size()[0]    accuracy = 100.0 * correct / total_samples    print('[xla:{}] Accuracy={:.2f}%'.format(        xm.get_ordinal(), accuracy), flush=True)    return accuracy, data, pred, target  # 训练和评估循环  accuracy = 0.0  data, pred, target = None, None, None  for epoch in range(1, FLAGS['num_epochs'] + 1):    para_loader = pl.ParallelLoader(train_loader, [device])    train_fun(para_loader.per_device_loader(device))    xm.master_print("Finished training epoch {}".format(epoch))    para_loader = pl.ParallelLoader(test_loader, [device])    accuracy, data, pred, target  = test_fun(para_loader.per_device_loader(device))    if FLAGS['metrics_debug']:      xm.master_print(met.metrics_report(), flush=True)  return accuracy, data, pred, target

当初，要将后果绘制为测试图像的预测标签和理论标签，将应用以下功能模块。

# 后果可视化import mathfrom matplotlib import pyplot as pltM, N = 5, 5RESULT_IMG_PATH = '/tmp/test_result.png'def plot_results(images, labels, preds):  images, labels, preds = images[:M*N], labels[:M*N], preds[:M*N]  inv_norm = transforms.Normalize((-0.1307/0.3081,), (1/0.3081,))  num_images = images.shape[0]  fig, axes = plt.subplots(M, N, figsize=(12, 12))  fig.suptitle('Predicted Lables')  for i, ax in enumerate(fig.axes):    ax.axis('off')    if i >= num_images:      continue    img, label, prediction = images[i], labels[i], preds[i]    img = inv_norm(img)    img = img.squeeze() # [1,Y,X] -> [Y,X]    label, prediction = label.item(), prediction.item()    if label == prediction:      ax.set_title(u'Actual {}/ Predicted {}'.format(label, prediction), color='blue')    else:      ax.set_title(          'Actual {}/ Predicted {}'.format(label, prediction), color='red')    ax.imshow(img)  plt.savefig(RESULT_IMG_PATH, transparent=True)

当初，咱们都筹备好在MNIST数据集上训练模型。训练开始前，咱们将记录开始工夫，训练完结后，咱们将记录完结工夫并打印50个epoch的总训练工夫。

# 启动训练流程def train_cnn(rank, flags):  global FLAGS  FLAGS = flags  torch.set_default_tensor_type('torch.FloatTensor')  accuracy, data, pred, target = train_mnist()  if rank == 0:    # 检索TPU外围0上的张量并绘制。    plot_results(data.cpu(), pred.cpu(), target.cpu())xmp.spawn(train_cnn, args=(FLAGS,), nprocs=FLAGS['num_cores'],          start_method='fork')

一旦训练胜利完结，咱们将打印训练所用的总工夫。

end_time = time.time()print('Total Training time = ',end_time-start_time )

正如咱们在下面看到的，这种办法破费了269秒或大概4.5分钟，这意味着50个epoch训练PyTorch模型不到5分钟。最初，咱们将通过训练的模型来可视化预测。

from google.colab.patches import cv2_imshowimport cv2img = cv2.imread(RESULT_IMG_PATH, cv2.IMREAD_UNCHANGED)cv2_imshow(img)

因而，咱们能够得出这样的论断：应用TPU实现深度学习模型能够实现疾速的训练，正如咱们后面所看到的那样。

在不到5分钟的工夫内，对50个epoch的40000张训练图像进行了CNN模型的训练。咱们在训练中也取得了89%以上的准确率。

因而，在TPU上训练深度学习模型在工夫和准确性方面总是有益处的。

参考文献：

  1. Joe Spisak, “Get started with PyTorch, Cloud TPUs, and Colab”.
  2. “PyTorch on XLA Devices”, PyTorch release.
  3. “Training PyTorch models on Cloud TPU Pods”, Google Cloud Guides.

原文链接：https://analyticsindiamag.com...

欢送关注磐创AI博客站：
http://panchuang.net/

sklearn机器学习中文官网文档：
http://sklearn123.com/

欢送关注磐创博客资源汇总站：
http://docs.panchuang.net/