使用TPU在PyTorch中实现ResNet50

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

PyTorch通过提供大量强大的工具和技术，一直在推动计算机视觉和深度学习领域的发展。

在计算机视觉领域，基于深度学习的执行需要处理大量的图像数据集，因此需要一个加速的环境来加快执行过程以达到可接受的精度水平。

PyTorch通过XLA（加速线性代数）提供了这一特性，XLA是一种线性代数编译器，可以针对多种类型的硬件，包括GPU和TPU。PyTorch/XLA环境与Google云TPU集成，实现了更快的执行速度。

在本文中，我们将在PyTorch中使用TPU演示一种深卷积神经网络ResNet50的实现。

该模型将在PyTorch/XLA环境中进行训练和测试，以完成CIFAR10数据集的分类任务。我们还将检查在50个epoch训练所花费的时间。

ResNet50在Pytorch的实现

为了利用TPU的功能，这个实现是在Google Colab中完成的。首先，我们需要从Notebook设置下的硬件加速器中选择TPU。

选择TPU后，我们将使用下面的行验证环境代码：

import os
assert 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)

在下一步中，我们将安装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

现在，我们将在这里导入所有必需的库。

from matplotlib import pyplot as plt
import numpy as np
import os
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.debug.metrics as met
import torch_xla.distributed.parallel_loader as pl
import torch_xla.distributed.xla_multiprocessing as xmp
import torch_xla.utils.utils as xu
import torchvision
from torchvision import datasets, transforms
import time
from google.colab.patches