PyTorch中文教程 | (1) PyTorch深度学习极速入门

GitHub 地址

本教程目标：

• 更高层次地理解PyTorch的Tensor库以及神经网络
• 训练一个小的神经网络模型用于图像分类

前提：安装torch和torchvision包。

目录

1. 什么是PyTorch？

2. Autograd：自动求导

3. 神经网络

4. 训练分类器

5. 可选：数据并行处理

---

1. 什么是PyTorch？

PyTorch是一个基于python的科学计算包，主要针对两类人群：

1) 作为NumPy的替代品，可以利用GPU的性能进行计算

2) 作为一个高灵活性、速度快的深度学习平台

• 入门

张量：

Tensor（张量）类似于NumPy的ndarray，但还可以在GPU上使用来加速计算。

from __future__ import print_function
import torch

创建一个没有初始化的5*3矩阵：

x = torch.empty(5, 3)
print(x)

创建一个随机初始化矩阵：

x = torch.rand(5, 3)
print(x)

构造一个填满0且数据类型为long的矩阵:

x = torch.zeros(5, 3, dtype=torch.long)
print(x)

直接从数据构造张量：

x = torch.tensor([5.5, 3])
print(x)

或者根据已有的tensor建立新的tensor。除非用户提供新的值，否则这些方法将重用输入张量的属性，例如dtype等：

x = x.new_ones(5, 3, dtype=torch.double)      # new_* methods take in sizes
print(x)

x = torch.randn_like(x, dtype=torch.float)    # 重载 dtype!
print(x)                                      # 结果有相同的size

获取tensor形状：

print(x.size())

注意：torch.Size本质上还是tuple，所以支持tuple的一切操作。

运算：

一种运算有多种语法。在下面的示例中，我们将研究加法运算（其他运算类似）。

加法：形式一

y = torch.rand(5, 3)
print(x + y)

加法：形式二：

print(torch.add(x, y))

加法：给定一个输出张量作为参数

result = torch.empty(5, 3)
torch.add(x, y, out=result)
print(result)

加法：原位/原地操作(in-place)

# adds x to y
y.add_(x)
print(y)

注意：任何一个in-place改变张量的操作后面都固定一个_。例如x.copy_(y)、x.t_()将更改x。

也可以使用像标准的NumPy一样的各种索引操作：

print(x)
print(x[:,1])#第2列

改变形状：如果想改变形状，可以使用torch.view

x = torch.randn(4, 4)
y = x.view(16)
z = x.view(-1, 8)  # the size -1 is inferred from other dimensions
print(x.size(), y.size(), z.size())

如果是仅包含一个元素的tensor，可以使用.item()来得到对应的python数值：

x = torch.randn(1)
print(x)
print(x.item())

超过100中tensor的运算操作，包括转置，索引，切片，数学运算， 线性代数，随机数等，具体访问这里。 

• NumPy桥

将一个Torch张量转换为一个NumPy数组是轻而易举的事情，反之亦然。

Torch张量和NumPy数组将共享它们的底层内存位置，更改一个将更改另一个。

将torch的Tensor转化为Numpy数组

a = torch.ones(5)
print(a)
b = a.numpy()
print(b)

看NumPy数组是如何改变里面的值的：

a.add_(1)
print(a)
print(b)

将NumPy数组转化为Torch张量：

看改变NumPy数组是如何自动改变Torch张量的：

import numpy as np

a = np.ones(5)
b = torch.from_numpy(a)
np.add(a,1,out=a)
print(a)
print(b)

CPU上的所有张量(CharTensor除外)都支持转换为NumPy以及由NumPy转换回来。

• CUDA上的张量

张量可以使用.to方法移动到任何设备（device）上：

# let us run this cell only if CUDA is available
# 我们将使用`torch.device`来将tensor移入和移出GPU
if torch.cuda.is_available():
    device = torch.device("cuda")          # a CUDA device object
    y = torch.ones_like(x, device=device)  # 直接在GPU上创建tensor
    x = x.to(device)                       # 或者使用`.to("cuda")`方法
    z = x + y
    print(z)
    print(z.to("cpu", torch.double))       # `.to`也能在移动时改变dtype

2. Autograd：自动求导

PyTorch中，所有神经网络的核心是autograd包。先简单介绍一下这个包，然后训练我们的第一个的神经网络。

autograd包为张量上的所有操作提供了自动求导机制。它是一个在运行时定义（define-by-run）的框架，这意味着反向传播是根据代码如何运行来决定的，并且每次迭代可以是不同的.

• 张量

torch.Tensor是这个包的核心类。如果设置它的属性 .requires_grad为True，那么它将会追踪对于该张量的所有操作。当完成计算后可以通过调用.backward()，来自动计算所有的梯度。这个张量的所有梯度将会自动累加到.grad属性。

要阻止一个张量被跟踪历史，可以调用.detach()方法将其与计算历史分离，并阻止它未来的计算记录被跟踪。

为了防止跟踪历史记录（和使用内存），可以将代码块包装在with torch.no_grad():中。在评估模型时特别有用，因为模型可能具有requires_grad = True的可训练的参数，但是我们不需要在此过程中对他们进行梯度计算。

还有一个类对于autograd的实现非常重要：Function。

Tensor和Function互相连接生成了一个非循环图，它编码了完整的计算历史。每个张量都有一个.grad_fn属性，它引用了一个创建了这个Tensor的Function（除非这个张量是用户手动创建的，即这个张量的grad_fn是None）。

如果需要计算导数，可以在Tensor上调用.backward()。如果Tensor是一个标量（即它包含一个元素的数据），则不需要为backward()指定任何参数，但是如果它有更多的元素，则需要指定一个gradient参数，它是形状匹配的张量。

创建一个张量并设置requires_grad=True用来追踪其计算历史

x = torch.ones(2, 2, requires_grad=True)
print(x)

y是计算的结果，所以它有grad_fn属性。

print(y.grad_fn)

z = y * y * 3
out = z.mean()

print(z, out)

.requires_grad_(...) 原地改变了现有张量的 requires_grad 标志。如果没有指定的话，默认输入的这个标志是False。

a = torch.randn(2, 2)
a = ((a * 3) / (a - 1))
print(a.requires_grad)
a.requires_grad_(True)
print(a.requires_grad)
b = (a * a).sum()
print(b.grad_fn)

• 梯度

因为out是一个标量。所以让我们直接进行反向传播，out.backward()和out.backward(torch.tensor(1.))等价。

out.backward()

输出导数d(out)/dx

print(x.grad)

得到都是4.5的矩阵。调用out张量 “o”。得到：

数学上，若有向量值函数y⃗ =f(x⃗ )，那么 y⃗ 相对于 x⃗ 的梯度是一个雅可比矩阵：

通常来说，torch.autograd 是计算雅可比向量积的一个“引擎”。也就是说，给定任意向量 :

，计算乘积 J⋅v。如果 v恰好是一个标量函数 l=g(y⃗ ) 的导数，即 

那么根据链式法则，雅可比向量积应该是 l对 x⃗  的导数:

雅可比向量积的这一特性使得将外部梯度输入到具有非标量输出的模型中变得非常方便。

现在我们来看一个雅可比向量积的例子:

x = torch.randn(3, requires_grad=True)

y = x * 2
while y.data.norm() < 1000:
    y = y * 2

print(y)

在这种情况下，y不再是标量。torch.autograd不能直接计算完整的雅可比矩阵，但是如果我们只想要雅可比向量积，只需将这个向量作为参数传给backward：

v = torch.tensor([0.1, 1.0, 0.0001], dtype=torch.float)
y.backward(v)

print(x.grad)

为了防止跟踪历史记录（和使用内存），可以将代码块包装在with torch.no_grad():中。在评估模型时特别有用，因为模型可能具有requires_grad = True的可训练的参数，但是我们不需要在此过程中对他们进行梯度计算。

也可以通过将代码块包装在 with torch.no_grad():中，来阻止autograd跟踪设置了 .requires_grad=True的张量的历史记录。

print(x.requires_grad)
print((x ** 2).requires_grad)

with torch.no_grad():
    print((x ** 2).requires_grad)

autograd And Function文档

3. 神经网络

可以使用torch.nn包来构建神经网络.

我们以及介绍了autograd，nn包依赖于autograd包来定义模型并对它们求导。一个nn.Module包含各个层和一个forward(input)方法，该方法返回output。

例如，下面这个神经网络可以对数字进行分类：

这是一个简单的前馈神经网络（feed-forward network）。它接受一个输入，然后将它送入下一层，一层接一层的传递，最后给出输出。

一个神经网络的典型训练过程如下：

1)定义包含一些可学习参数（或者叫权重）的神经网络

2)在输入数据集上迭代

3)通过网络处理输入

4)计算损失（输出和正确答案的距离）

5)将梯度反向传播给网络的参数

6)更新网络的权重，一般使用一个简单的规则：weight = weight - learning_rate * gradient

• 定义网络

import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        # 输入图像channel：1；输出channel：6；5x5卷积核
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # 2x2 Max pooling
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # If the size is a square you can only specify a single number
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # 除去批大小维度的其余维度
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


net = Net()
print(net)

我们只需要定义 forward 函数，backward函数会在使用autograd时自动定义，backward函数用来计算导数。可以在 forward 函数中使用任何针对张量的操作和计算。

一个模型的可学习参数可以通过net.parameters()返回：

params = list(net.parameters())
print(len(params))
print(params[0].size())  # conv1's .weight

让我们尝试一个随机的32x32的输入。注意，这个网络（LeNet）的期待输入是32x32。如果使用MNIST数据集来训练这个网络，要把图片大小重新调整到32x32。

input = torch.randn(1, 1, 32, 32)
out = net(input)
print(out)

清零所有参数的梯度缓存，然后进行随机梯度的反向传播：

net.zero_grad()
out.backward(torch.randn(1, 10))

注意：torch.nn只支持小批量处理（mini-batches）。整个torch.nn包只支持小批量样本的输入，不支持单个样本。比如，nn.Conv2d 接受一个4维的张量，即nSamples x nChannels x Height x Width。如果是一个单独的样本，只需要用input.unsqueeze(0)来添加一个“假的”批大小维度。

复习：

1）torch.Tensor - 一个多维数组，支持诸如backward()等的自动求导操作，同时也保存了张量的梯度。

2）nn.Module - 神经网络模块。是一种方便封装参数的方式，具有将参数移动到GPU、导出、加载等功能。

3）nn.Parameter - 张量的一种，当它作为一个属性分配给一个Module时，它会被自动注册为一个参数。

4）autograd.Function - 实现了自动求导前向和反向传播的定义，每个Tensor至少创建一个Function节点，该节点连接到创建Tensor的函数并对其历史进行编码。

至今我们讨论了：

1）定义一个神经网络

2）处理输入调用backard

还剩：

1）计算损失

2）更新网络权重

• 损失函数

一个损失函数接受一对(output, target)作为输入，计算一个值来估计网络的输出和目标值相差多少。

nn包中有很多不同的损失函数。nn.MSELoss是比较简单的一种，它计算输出和目标的均方误差（mean-squared error）。

output = net(input)
target = torch.randn(10)  # a dummy target, for example
target = target.view(1, -1)  # make it the same shape as output 行向量(2维tensor表示)
criterion = nn.MSELoss()

loss = criterion(output, target)
print(loss)

现在，如果使用loss的.grad_fn属性跟踪反向传播过程，会看到计算图如下：

input -> conv2d -> relu -> maxpool2d -> conv2d -> relu -> maxpool2d
      -> view -> linear -> relu -> linear -> relu -> linear
      -> MSELoss
      -> loss

所以，当我们调用loss.backward()，整张图开始关于loss微分，图中所有设置了requires_grad=True的张量的.grad属性累积着梯度张量。

为了说明这一点，让我们向后跟踪几步：

print(loss.grad_fn)  # MSELoss
print(loss.grad_fn.next_functions[0][0])  # Linear
print(loss.grad_fn.next_functions[0][0].next_functions[0][0])  # ReLU

• 反向传播

我们只需要调用loss.backward()来反向传播权重。我们需要清零现有的梯度，否则梯度将会与已有的梯度累加。

现在，我们将调用loss.backward()，并查看conv1层的偏置（bias）在反向传播前后的梯度。

net.zero_grad()     # 清零所有参数（parameter）的梯度缓存

print('conv1.bias.grad before backward')
print(net.conv1.bias.grad)

loss.backward()

print('conv1.bias.grad after backward')
print(net.conv1.bias.grad)

神经网络包包含了各种模块和损失函数，这些模块和损失函数构成了深度神经网络的构建模块。完整的文档列表见这里。

• 更新网络权重

最简单的更新规则是随机梯度下降法（SGD）:

我们可以使用简单的python代码来实现:

learning_rate = 0.01
for f in net.parameters():
    f.data.sub_(f.grad.data * learning_rate)

然而，在使用神经网络时，可能希望使用各种不同的更新规则，如SGD、Nesterov-SGD、Adam、RMSProp等。为此，我们构建了一个较小的包torch.optim，它实现了所有的这些方法。使用它很简单：

import torch.optim as optim

# 创建优化器（optimizer）
optimizer = optim.SGD(net.parameters(), lr=0.01)

# 在训练的迭代中：
optimizer.zero_grad()   # 清零梯度缓存
output = net(input)
loss = criterion(output, target)
loss.backward()
optimizer.step()    # 更新参数

观察梯度缓存区是如何使用optimizer.zero_grad()手动清零的。这是因为梯度是累加的，正如前面反向传播章节叙述的那样。

4. 训练分类器

目前为止，我们看到了如何定义网络，计算损失，并更新网络的权重。

• 数据

通常来说，当必须处理图像、文本、音频或视频数据时，可以使用python标准库将数据加载到numpy数组里。然后将这个数组转化成torch.*Tensor。

1）对于图片，有Pillow，OpenCV等包可以使用

2）对于音频，有scipy和librosa等包可以使用

3）对于文本，不管是原生python的或者是基于Cython的文本，可以使用NLTK和SpaCy

特别对于视觉方面，我们创建了一个包，名字叫torchvision，其中包含了针对Imagenet、CIFAR10、MNIST等常用数据集的数据加载器（data loaders），还有对图片数据变形的操作，即torchvision.datasets和torch.utils.data.DataLoader。

这提供了极大的便利，可以避免编写样板代码。

在这个教程中，我们将使用CIFAR10数据集，它有如下的分类：“飞机”，“汽车”，“鸟”，“猫”，“鹿”，“狗”，“青蛙”，“马”，“船”，“卡车”等。在CIFAR-10里面的图片数据大小是3x32x32，即三通道彩色图，图片大小是32x32像素。

• 训练一个图像分类器

我们将按顺序做以下步骤：

1）通过torchvision加载CIFAR10里面的训练和测试数据集，并对数据进行标准化

2）定义卷积神经网络

3）定义损失函数

4）利用训练数据训练网络

5）利用测试数据测试网络

加载并标准化CIFAR-10

使用torchvision加载CIFAR10超级简单。

import torch
import torchvision
import torchvision.transforms as transforms

torchvision数据集加载完后的输出是范围在[0, 1]之间的PILImage。我们将其标准化为范围在[-1, 1]之间的张量。

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

我们可以展示一些训练数据图片：

import matplotlib.pyplot as plt
import numpy as np

# 输出图像的函数


def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


# 随机得到一些训练图片
dataiter = iter(trainloader)
images, labels = dataiter.next()

# 显示图片
imshow(torchvision.utils.make_grid(images))
# 打印图片标签
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

 

• 定义卷积神经网络

和之前的网络几乎一样，只是要把单通道改为3通道，彩色图：

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()

(32,32,3) ->conv1 5*5*3 6 stride:（1，1） 不填充->(28,28,6)->pool (2,2)  stride:（1，1）->(14,14,6) ->conv2 5*5*6 16  stride:（1，1） 不填充->(10,10,16)->pool (2,2)  stride:（1，1）->(5,5,16)->view->fc...

• 定义损失函数和优化器

我们使用分类的交叉熵损失和随机梯度下降（使用momentum）。

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

• 训练网络

我们只需要遍历我们的数据迭代器，并且输入“喂”给网络和优化函数。

for epoch in range(2):  # loop over the dataset multiple times

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

print('Finished Training')

 

• 使用测试数据测试网络

我们已经在训练集上训练了2遍网络。但是我们需要检查网络是否学到了一些东西。

我们将通过预测神经网络输出的标签来检查这个问题，并和正确样本进行（ground-truth）对比。如果预测是正确的，我们将样本添加到正确预测的列表中。

ok，第一步。让我们显示测试集中的图像来熟悉一下。

dataiter = iter(testloader)
images, labels = dataiter.next()

# 输出图片
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))

 

ok，现在让我们看看神经网络认为上面的例子是:

outputs = net(images)

输出是10个类别的量值。一个类的值越高，网络就越认为这个图像属于这个特定的类。让我们得到最高量值的下标/索引；

_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join('%5s' % classes[predicted[j]] for j in range(4)))

 

让我们看看网络在整个测试集上表现的怎么样。

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))

这比随机选取（即从10个类中随机选择一个类，正确率是10%）要好很多。看来网络确实学到了一些东西。那么哪些是表现好的类呢？哪些是表现的差的类呢？

class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1


for i in range(10):
    print('Accuracy of %5s : %2d %%' % (
        classes[i], 100 * class_correct[i] / class_total[i]))

• 在GPU上训练

与将一个张量传递给GPU一样，可以这样将神经网络转移到GPU上。

如果我们有cuda可用的话，让我们首先定义第一个设备为可见cuda设备：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Assuming that we are on a CUDA machine, this should print a CUDA device:

print(device)

本节的其余部分假设device是CUDA。

然后这些方法将递归遍历所有模块，并将它们的参数和缓冲区转换为CUDA张量：

net.to(device)

请记住，我们不得不将输入和目标在每一步都送入GPU：

inputs, labels = inputs.to(device), labels.to(device)

为什么我们感受不到与CPU相比的巨大加速？因为我们的网络实在是太小了。

尝试一下：加宽你的网络（注意第一个nn.Conv2d的第二个参数（filter个数）和第二个nn.Conv2d的第一个参数（通道）要相同），看看能获得多少加速。

已实现的目标：

1）在更高层次上理解PyTorch的Tensor库和神经网络

2）训练一个小的神经网络做图片分类

接下来可以尝试：

Train neural nets to play video games

Train a state-of-the-art ResNet network on imagenet

Train a face generator using Generative Adversarial Networks

Train a word-level language model using Recurrent LSTM networks

More examples

More tutorials

Discuss PyTorch on the Forums

Chat with other users on Slack

5. 可选：数据并行处理

我们将学习如何使用数据并行（DataParallel）来使用多GPU。

PyTorch非常容易的就可以使用GPU，可以用如下方式把一个模型放到GPU上：

device = torch.device("cuda：0")
model.to(device)

然后可以复制所有的张量到GPU上：

mytensor = my_tensor.to(device)

请注意，调用my_tensor.to(device)返回一个GPU上的my_tensor副本，而不是重写my_tensor。我们需要把它赋值给一个新的张量并在GPU上使用这个张量。

在多GPU上执行前向和反向传播是自然而然的事。然而，PyTorch默认将只是用一个GPU。你可以使用DataParallel让模型并行运行来轻易的让你的操作在多个GPU上运行。

model = nn.DataParallel(model)

• 导入和参数

导入PyTorch模块和定义参数。

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# Parameters 和 DataLoaders
input_size = 5
output_size = 2

batch_size = 30
data_size = 100

设备（Device）：

device = torch.device("cuda：0" if torch.cuda.is_available() else "cpu")

• 虚拟数据集

要制作一个虚拟（随机）数据集，只需实现__getitem__。

class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len

rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
                         batch_size=batch_size, shuffle=True)

• 简单模型

作为演示，我们的模型只接受一个输入，执行一个线性操作，然后得到结果。然而，你能在任何模型（CNN，RNN，Capsule Net等）上使用DataParallel。

我们在模型内部放置了一条打印语句来检测输入和输出向量的大小。请注意批等级为0时打印的内容。

class Model(nn.Module):
    # Our model

    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model： input size", input.size(),
              "output size", output.size())

        return output

• 创建一个模型和数据并行

首先，我们需要创建一个模型实例和检测我们是否有多个GPU。如果我们有多个GPU，我们使用nn.DataParallel来包装我们的模型。然后通过model.to(device)把模型放到GPU上。

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1：
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  model = nn.DataParallel(model)

model.to(device)

• 运行模型

for data in rand_loader：
    input = data.to(device)
    output = model(input)
    print("Outside： input size", input.size(),
          "output_size", output.size())

In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
        In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
        In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
        In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
        In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
        In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
        In Model： input size torch.Size([5, 5]) output size torch.Size([5, 2])
        In Model： input size torch.Size([5, 5]) output size torch.Size([5, 2])
Outside： input size torch.Size([10, 5]) output_size torch.Size([10, 2])

• 结果

当我们对30个输入和输出进行批处理时，我们和期望的一样得到30个输入和30个输出，但是若有多个GPU，会得到如下的结果。

2个GPU：

Let's use 2 GPUs!
    In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
    In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
    In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
    In Model： input size torch.Size([15, 5]) output size torch.Size([15, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model： input size torch.Size([5, 5]) output size torch.Size([5, 2])
    In Model： input size torch.Size([5, 5]) output size torch.Size([5, 2])
Outside： input size torch.Size([10, 5]) output_size torch.Size([10, 2])

3个GPU：

Let's use 3 GPUs!
    In Model： input size torch.Size([10, 5]) output size torch.Size([10, 2])
    In Model： input size torch.Size([10, 5]) output size torch.Size([10, 2])
    In Model： input size torch.Size([10, 5]) output size torch.Size([10, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model： input size torch.Size([10, 5]) output size torch.Size([10, 2])
    In Model： input size torch.Size([10, 5]) output size torch.Size([10, 2])
    In Model： input size torch.Size([10, 5]) output size torch.Size([10, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model： input size torch.Size([10, 5]) output size torch.Size([10, 2])
    In Model： input size torch.Size([10, 5]) output size torch.Size([10, 2])
    In Model： input size torch.Size([10, 5]) output size torch.Size([10, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([2, 5]) output size torch.Size([2, 2])
Outside： input size torch.Size([10, 5]) output_size torch.Size([10, 2])

8个GPU：

Let's use 8 GPUs!
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([2, 5]) output size torch.Size([2, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([2, 5]) output size torch.Size([2, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([4, 5]) output size torch.Size([4, 2])
    In Model： input size torch.Size([2, 5]) output size torch.Size([2, 2])
Outside： input size torch.Size([30, 5]) output_size torch.Size([30, 2])
    In Model： input size torch.Size([2, 5]) output size torch.Size([2, 2])
    In Model： input size torch.Size([2, 5]) output size torch.Size([2, 2])
    In Model： input size torch.Size([2, 5]) output size torch.Size([2, 2])
    In Model： input size torch.Size([2, 5]) output size torch.Size([2, 2])
    In Model： input size torch.Size([2, 5]) output size torch.Size([2, 2])
Outside： input size torch.Size([10, 5]) output_size torch.Size([10, 2])

• 总结

DataParallel自动的划分数据，并将作业发送到多个GPU上的多个模型。在每个模型完成作业后，DataParallel收集并合并结果返回给你。

更多信息