正文 

当前位置： 首页 > 程序 > java教程

pytorch实现tpu版本cnn模型(代码片段)

panchuangai panchuangai     2022-12-09     763

关键词：

作者|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 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)

技术图片

启用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 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
from torchvision import datasets, transforms

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

# 定义参数
FLAGS = 
FLAGS[‘datadir‘] = "/tmp/mnist"
FLAGS[‘batch_size‘] = 128
FLAGS[‘num_workers‘] = 4
FLAGS[‘learning_rate‘] = 0.01
FLAGS[‘momentum‘] = 0.5
FLAGS[‘num_epochs‘] = 50
FLAGS[‘num_cores‘] = 8
FLAGS[‘log_steps‘] = 20
FLAGS[‘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 math
from matplotlib import pyplot as plt

M, N = 5, 5
RESULT_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_imshow
import cv2
img = 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/how-to-implement-cnn-model-using-pytorch-with-tpu/

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

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

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






基于pytorch使用实现cnn如何使用pytorch构建cnn卷积神经网络(代码片段)

基于pytorch使用实现CNN如何使用pytorch构建CNN卷积神经网络本文是一个基于pytorch使用CNN在生物信息学上进行位点预测的例子基于pytorch实现CNN，基于CNN进行位点预测，将CNN代码进行封装，可以非常简单的使用代码，基... 查看详情

使用 pytorch-lightning 实现 Network in Network CNN 模型

】使用pytorch-lightning实现NetworkinNetworkCNN模型【英文标题】：ImplementaNetworkinNetworkCNNmodelusingpytorch-lightning【发布时间】：2021-08-2721:32:06【问题描述】：我正在尝试实现NiN模型。基本上试图从d2l复制代码这是我的代码。importpandasaspdimp... 查看详情

ai常用框架和工具丨13.pytorch实现基于cnn的手写数字识别

代码实例，PyTorch实现基于CNN的手写数字识别，希望对您有所帮助。文章目录环境说明一、模型训练1.1导入相关依赖1.2选择使用的硬件1.3超参数配置1.4准备训练集和测试集1.5模型创建1.6模型评估指标1.7模型训练1.8模型测试1.9模型... 查看详情

ai常用框架和工具丨13.pytorch实现基于cnn的手写数字识别

代码实例，PyTorch实现基于CNN的手写数字识别，希望对您有所帮助。文章目录环境说明一、模型训练1.1导入相关依赖1.2选择使用的硬件1.3超参数配置1.4准备训练集和测试集1.5模型创建1.6模型评估指标1.7模型训练1.8模型测试1.9模型... 查看详情

ai常用框架和工具丨13.pytorch实现基于cnn的手写数字识别

代码实例，PyTorch实现基于CNN的手写数字识别，希望对您有所帮助。文章目录环境说明一、模型训练1.1导入相关依赖1.2选择使用的硬件1.3超参数配置1.4准备训练集和测试集1.5模型创建1.6模型评估指标1.7模型训练1.8模型测试1.9模型... 查看详情

基于pytorch使用实现cnn如何使用pytorch构建cnn卷积神经网络(代码片段)

基于pytorch使用实现CNN如何使用pytorch构建CNN卷积神经网络本文是一个基于pytorch使用CNN在生物信息学上进行位点预测的例子基于pytorch实现CNN，基于CNN进行位点预测，将CNN代码进行封装，可以非常简单的使用代码，基... 查看详情

pytorch实战笔记——cnn实现情感分析(代码片段)

本文展示的是使用Pytorch构建一个TextCNN来实现情感分析。本文的架构是第一章详细介绍TextCNN（不带公式版），第二章是核心代码部分。目录1.TextCNN2.TextCNN实现情感分析参考1.TextCNN相较于LSTM而言，我个人其实是没看... 查看详情

在 Colab TPU 上运行 Pytorch 堆叠模型

】在ColabTPU上运行Pytorch堆叠模型【英文标题】：RunPytorchstackedmodelonColabTPU【发布时间】：2021-03-1423:41:15【问题描述】：我正在尝试在Colab多核TPU上运行我的模型，但我真的不知道该怎么做。我尝试了thistutorialnotebook，但出现了一些... 查看详情

5.使用pytorch预先训练的模型执行目标检测(代码片段)

5.使用PyTorch预先训练的网络执行目标检测这篇博客将介绍如何使用PyTorch预训练的网络执行目标检测，这些网络是开创性的、最先进的图像分类网络，包括使用ResNet的更快R-CNN、使用MobileNet的更快R-CNN和RetinaNet。具有ResNet50... 查看详情

用rnn&cnn进行情感分析-pytorch(代码片段)

动手学深度学习笔记一、情感分析及数据集1.读取数据集2.预处理数据集二、利用RNN进行情感分析1.使用RNN表示单个文本2.加载预训练的词向量3.训练和评估模型4.预测三、利用CNN进行情感分析1.一维卷积2.textCNN模型3.加载预训练词... 查看详情

cnn经典网络模型：resnet简介及代码实现（pytorch超详细注释版）(代码片段)

目录一、开发背景二、网络结构三、模型特点四、代码实现1.model.py2.train.py3.predict.py4.spilit_data.py五、参考内容一、开发背景残差神经网络(ResNet)是由微软研究院的何恺明、张祥雨、任少卿、孙剑等人提出的，斩获2015年ImageNet... 查看详情

将 CNN 模型代码从 Keras 转换为 Pytorch

】将CNN模型代码从Keras转换为Pytorch【英文标题】：ToconvertCNNmodelcodefromKerastoPytorch【发布时间】：2022-01-0218:37:12【问题描述】：我正在尝试。这里是KerasSequential层model=Sequential()model.add(Conv2D(filters=64,kernel_size=(3,3),activation="relu",input_sh... 查看详情

图像分类convit从入门到实战——使用convit实现植物幼苗的分类（pytorch）(代码片段)

摘要来自Facebook的研究者提出了一种名为ConViT的新计算机视觉模型，它结合了两种广泛使用的AI架构——卷积神经网络(CNN)和Transformer，该模型取长补短，克服了CNN和Transformer本身的一些局限性。同时，借助这两种架... 查看详情

小白学习pytorch教程九基于pytorch训练第一个rnn模型(代码片段)

@Author：Runsen当阅读一篇课文时，我们可以根据前面的单词来理解每个单词的，而不是从零开始理解每个单词。这可以称为记忆。卷积神经网络模型（CNN）不能实现这种记忆，因此引入了递归神经网络模型（RNN）来解决这一问题... 查看详情

[pytorch系列-61]：循环神经网络-中文新闻文本分类详解-3-cnn网络训练与评估代码详解(代码片段)

...6.1测试方法的定义6.2开始测试第1章预备知识1.1业务概述[Pytorch系列-59]：循环神经网络-中文新闻文本分类详解-1-业务目标分析与总体架构_文火冰糖（王文兵）的博客-CSDN博客https://blog.csdn.net/HiWangWenBing/article/details/121756... 查看详情

pytorch实现卷积神经网络cnn

Pytorch是torch的Python版本，对TensorFlow造成很大的冲击，TensorFlow无疑是最流行的，但是Pytorch号称在诸多性能上要优于TensorFlow，比如在RNN的训练上，所以Pytorch也吸引了很多人的关注。之前有一篇关于TensorFlow实现的CNN可以用来做对比... 查看详情

pytorch从头搭建并训练一个神经网络模型（图像分类cnn）(代码片段)

...多稍微修改下源码的接口满足自己的需求。还从来没有用PyTorch从头搭建并训练一个模型出来。正好最近在较为系统地学PyTorch，就总结一下如何从头搭建并训练一个神经网络模型。1.使用torchvision加载数据集并做预处理我们使... 查看详情

pytorch版本的bert模型代码(代码片段)

代码出处：https://github.com/aespresso/a_journey_into_math_of_ml/blob/master/04_transformer_tutorial_2nd_part/BERT_tutorial/models/bert_model.py#coding=utf-8#Copyright2018TheGoogleAILanguageTeamAuthorsandT 查看详情