手把手教你由TensorFlow上手PyTorch（附代码）

当我第一次尝试学习 PyTorch 时，没几天就放弃了。和 TensorFlow 相比，我很难弄清 PyTorch 的核心要领。但是随后不久，PyTorch 发布了一个新版本，我决定重新来过。在第二次的学习中，我开始了解这个框架的易用性。在本文中，我会简要解释 PyTorch 的核心概念，为你转入这个框架提供一些必要的动力。其中包含了一些基础概念，以及先进的功能如学习速率调整、自定义层等等。

PyTorch 的易用性如何？Andrej Karpathy 是这样评价的

资源

首先要知道的是：PyTorch 的主目录和教程是分开的。而且因为开发和版本更新的速度过快，有时候两者之间并不匹配。所以你需要不时查看源代码：http://pytorch.org/tutorials/。

当然，目前网络上已有了一些 PyTorch 论坛，你可以在其中询问相关的问题，并很快得到回复：https://discuss.pytorch.org/。

把 PyTorch 当做 NumPy 用

让我们先看看 PyTorch 本身，其主要构件是张量——这和 NumPy 看起来差不多。这种性质使得 PyTorch 可支持大量相同的 API，所以有时候你可以把它用作是 NumPy 的替代品。PyTorch 的开发者们这么做的原因是希望这种框架可以完全获得 GPU 加速带来的便利，以便你可以快速进行数据预处理，或其他任何机器学习任务。将张量从 NumPy 转换至 PyTorch 非常容易，反之亦然。让我们看看如下代码：

1. import torch
2. import numpy as np
4. numpy_tensor = np.random.randn(10, 20)
6. # convert numpy array to pytorch array
7. pytorch_tensor = torch.Tensor(numpy_tensor)
8. # or another way
9. pytorch_tensor = torch.from_numpy(numpy_tensor)
11. # convert torch tensor to numpy representation
12. pytorch_tensor.numpy()
14. # if we want to use tensor on GPU provide another type
15. dtype = torch.cuda.FloatTensor
16. gpu_tensor = torch.randn(10, 20).type(dtype)
17. # or just call `cuda()` method
18. gpu_tensor = pytorch_tensor.cuda()
19. # call back to the CPU
20. cpu_tensor = gpu_tensor.cpu()
22. # define pytorch tensors
23. x = torch.randn(10, 20)
24. y = torch.ones(20, 5)
25. # `@` mean matrix multiplication from python3.5, PEP-0465
26. res = x @ y
28. # get the shape
29. res.shape  # torch.Size([10, 5])

从张量到变量

张量是 PyTorch 的一个完美组件，但是要想构建神经网络这还远远不够。反向传播怎么办？当然，我们可以手动实现它，但是真的需要这样做吗？幸好还有自动微分。为了支持这个功能，PyTorch 提供了变量，它是张量之上的封装。如此，我们可以构建自己的计算图，并自动计算梯度。每个变量实例都有两个属性：包含初始张量本身的.data，以及包含相应张量梯度的.grad

1. import torch
2. from torch.autograd import Variable
4. # define an inputs
5. x_tensor = torch.randn(10, 20)
6. y_tensor = torch.randn(10, 5)
7. x = Variable(x_tensor, requires_grad=False)
8. y = Variable(y_tensor, requires_grad=False)
9. # define some weights
10. w = Variable(torch.randn(20, 5), requires_grad=True)
12. # get variable tensor
13. print(type(w.data))  # torch.FloatTensor
14. # get variable gradient
15. print(w.grad)  # None
17. loss = torch.mean((y - x @ w) ** 2)
19. # calculate the gradients
20. loss.backward()
21. print(w.grad)  # some gradients
22. # manually apply gradients
23. w.data -= 0.01 * w.grad.data
24. # manually zero gradients after update
25. w.grad.data.zero_()

你也许注意到我们手动计算了自己的梯度，这样看起来很麻烦，我们能使用优化器吗？当然。

1. import torch
2. from torch.autograd import Variable
3. import torch.nn.functional as F
6. x = Variable(torch.randn(10, 20), requires_grad=False)
7. y = Variable(torch.randn(10, 3), requires_grad=False)
8. # define some weights
9. w1 = Variable(torch.randn(20, 5), requires_grad=True)
10. w2 = Variable(torch.randn(5, 3), requires_grad=True)
12. learning_rate = 0.1
13. loss_fn = torch.nn.MSELoss()
14. optimizer = torch.optim.SGD([w1, w2], lr=learning_rate)
15. for step in range(5):
16.    pred = F.sigmoid(x @ w1)
17.    pred = F.sigmoid(pred @ w2)
18.    loss = loss_fn(pred, y)
20.    # manually zero all previous gradients
21.    optimizer.zero_grad()
22.    # calculate new gradients
23.    loss.backward()
24.    # apply new gradients
25.    optimizer.step()

并不是所有的变量都可以自动更新。但是你应该可以从最后一段代码中看到重点：我们仍然需要在计算新梯度之前将它手动归零。这是 PyTorch 的核心理念之一。有时我们会不太明白为什么要这么做，但另一方面，这样可以让我们充分控制自己的梯度。

静态图 vs 动态图

PyTorch 和 TensorFlow 的另一个主要区别在于其不同的计算图表现形式。TensorFlow 使用静态图，这意味着我们是先定义，然后不断使用它。在 PyTorch 中，每次正向传播都会定义一个新计算图。在开始阶段，两者之间或许差别不是很大，但动态图会在你希望调试代码，或定义一些条件语句时显现出自己的优势。就像你可以使用自己最喜欢的 debugger 一样！

你可以比较一下 while 循环语句的下两种定义——第一个是 TensorFlow 中，第二个是 PyTorch 中：

1. import tensorflow as tf
4. first_counter = tf.constant(0)
5. second_counter = tf.constant(10)
6. some_value = tf.Variable(15)
9. # condition should handle all args:
10. def cond(first_counter, second_counter, *args):
11.    return first_counter < second_counter
14. def body(first_counter, second_counter, some_value):
15.    first_counter = tf.add(first_counter, 2)
16.    second_counter = tf.add(second_counter, 1)
17.    return first_counter, second_counter, some_value
20. c1, c2, val = tf.while_loop(
21.    cond, body, [first_counter, second_counter, some_value])
23. with tf.Session() as sess:
24.    sess.run(tf.global_variables_initializer())
25.    counter_1_res, counter_2_res = sess.run([c1, c2])

1. import torch
3. first_counter = torch.Tensor([0])
4. second_counter = torch.Tensor([10])
5. some_value = torch.Tensor(15)
7. while (first_counter < second_counter)[0]:
8.    first_counter += 2
9.    second_counter += 1

看起来第二种方法比第一个简单多了，你觉得呢？

模型定义

现在我们看到，想在 PyTorch 中创建 if/else/while 复杂语句非常容易。不过让我们先回到常见模型中，PyTorch 提供了非常类似于 Keras 的、即开即用的层构造函数：

神经网络包（nn）定义了一系列的模块，它可以粗略地等价于神经网络的层。模块接收输入变量并计算输出变量，但也可以保存内部状态，例如包含可学习参数的变量。nn 包还定义了一组在训练神经网络时常用的损失函数。

1. from collections import OrderedDict
3. import torch.nn as nn
6. # Example of using Sequential
7. model = nn.Sequential(
8.    nn.Conv2d(1, 20, 5),
9.    nn.ReLU(),
10.    nn.Conv2d(20, 64, 5),
11.    nn.ReLU()
12. )
14. # Example of using Sequential with OrderedDict
15. model = nn.Sequential(OrderedDict([
16.    ('conv1', nn.Conv2d(1, 20, 5)),
17.    ('relu1', nn.ReLU()),
18.    ('conv2', nn.Conv2d(20, 64, 5)),
19.    ('relu2', nn.ReLU())
20. ]))
22. output = model(some_input)

如果你想要构建复杂的模型，我们可以将 nn.Module 类子类化。当然，这两种方式也可以互相结合。

1. from torch import nn
4. class Model(nn.Module):
5.    def __init__(self):
6.        super().__init__()
7.        self.feature_extractor = nn.Sequential(
8.            nn.Conv2d(3, 12, kernel_size=3, padding=1, stride=1),
9.            nn.Conv2d(12, 24, kernel_size=3, padding=1, stride=1),
10.        )
11.        self.second_extractor = nn.Conv2d(
12.            24, 36, kernel_size=3, padding=1, stride=1)
14.    def forward(self, x):
15.        x = self.feature_extractor(x)
16.        x = self.second_extractor(x)
17.        # note that we may call same layer twice or mode
18.        x = self.second_extractor(x)
19.        return x

在__init__方法中，我们需要定义之后需要使用的所有层。在正向方法中，我们需要提出如何使用已经定义的层的步骤。而在反向传播上，和往常一样，计算是自动进行的。

自定义层

如果我们想要定义一些非标准反向传播模型要怎么办？这里有一个例子——XNOR 网络：

在这里我们不会深入细节，如果你对它感兴趣，可以参考一下原始论文：https://arxiv.org/abs/1603.05279

与我们问题相关的是反向传播需要权重必须介于-1 到 1 之间。在 PyTorch 中，这可以很容易实现：

1. import torch
4. class MyFunction(torch.autograd.Function):
6.    @staticmethod
7.    def forward(ctx, input):
8.        ctx.save_for_backward(input)
9.        output = torch.sign(input)
10.        return output
12.    @staticmethod
13.    def backward(ctx, grad_output):
14.        # saved tensors - tuple of tensors, so we need get first
15.        input, = ctx.saved_variables
16.        grad_output[input.ge(1)] = 0
17.        grad_output[input.le(-1)] = 0
18.        return grad_output
21. # usage
22. x = torch.randn(10, 20)
23. y = MyFunction.apply(x)
24. # or
25. my_func = MyFunction.apply
26. y = my_func(x)
29. # and if we want to use inside nn.Module
30. class MyFunctionModule(torch.nn.Module):
31.    def forward(self, x):
32.        return MyFunction.apply(x)

正如你所见，我们应该只定义两种方法：一个为正向传播，一个为反向传播。如果我们需要从正向通道访问一些变量，我们可以将它们存储在 ctx 变量中。注意：在此前的 API 正向/反向传播不是静态的，我们存储变量需要以 self.save_for_backward(input) 的形式，并以 input, _ = self.saved_tensors 的方式接入。

在 CUDA 上训练模型

我们曾经讨论过传递一个张量到 CUDA 上。但如果希望传递整个模型，我们可以通过调用.cuda() 来完成，并将每个输入变量传递到.cuda() 中。在所有计算后，我们需要用返回.cpu() 的方法来获得结果。

同时，PyTorch 也支持在源代码中直接分配设备

1. import torch
3. ### tensor example
4. x_cpu = torch.randn(10, 20)
5. w_cpu = torch.randn(20, 10)
6. # direct transfer to the GPU
7. x_gpu = x_cpu.cuda()
8. w_gpu = w_cpu.cuda()
9. result_gpu = x_gpu @ w_gpu
10. # get back from GPU to CPU
11. result_cpu = result_gpu.cpu()
13. ### model example
14. model = model.cuda()
15. # train step
16. inputs = Variable(inputs.cuda())
17. outputs = model(inputs)
18. # get back from GPU to CPU
19. outputs = outputs.cpu()

因为有些时候我们想在 CPU 和 GPU 中运行相同的模型，而无需改动代码，我们会需要一种封装：

1. class Trainer:
2.    def __init__(self, model, use_cuda=False, gpu_idx=0):
3.        self.use_cuda = use_cuda
4.        self.gpu_idx = gpu_idx
5.        self.model = self.to_gpu(model)
7.    def to_gpu(self, tensor):
8.        if self.use_cuda:
9.            return tensor.cuda(self.gpu_idx)
10.        else:
11.            return tensor
13.    def from_gpu(self, tensor):
14.        if self.use_cuda:
15.            return tensor.cpu()
16.        else:
17.            return tensor
19.    def train(self, inputs):
20.        inputs = self.to_gpu(inputs)
21.        outputs = self.model(inputs)
22.        outputs = self.from_gpu(outputs)

权重初始化

在 TesnorFlow 中权重初始化主要是在张量声明中进行的。PyTorch 则提供了另一种方法：首先声明张量，随后在下一步里改变张量的权重。权重可以用调用 torch.nn.init 包中的多种方法初始化为直接访问张量的属性。这个决定或许并不直接了当，但当你希望初始化具有某些相同初始化类型的层时，它就会变得有用。

1. import torch
2. from torch.autograd import Variable
4. # new way with `init` module
5. w = torch.Tensor(3, 5)
6. torch.nn.init.normal(w)
7. # work for Variables also
8. w2 = Variable(w)
9. torch.nn.init.normal(w2)
10. # old styled direct access to tensors data attribute
11. w2.data.normal_()
13. # example for some module
14. def weights_init(m):
15.    classname = m.__class__.__name__
16.    if classname.find('Conv') != -1:
17.        m.weight.data.normal_(0.0, 0.02)
18.    elif classname.find('BatchNorm') != -1:
19.        m.weight.data.normal_(1.0, 0.02)
20.        m.bias.data.fill_(0)
22. # for loop approach with direct access
23. class MyModel(nn.Module):
24.    def __init__(self):
25.        for m in self.modules():
26.            if isinstance(m, nn.Conv2d):
27.                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
28.                m.weight.data.normal_(0, math.sqrt(2. / n))
29.            elif isinstance(m, nn.BatchNorm2d):
30.                m.weight.data.fill_(1)
31.                m.bias.data.zero_()
32.            elif isinstance(m, nn.Linear):
33.                m.bias.data.zero_()

反向排除子图

有时，当你希望保留模型中的某些层或者为生产环境做准备的时候，禁用某些层的自动梯度机制非常有用。在这种思路下，PyTorch 设计了两个 flag：requires_grad 和 volatile。第一个可以禁用当前层的梯度，但子节点仍然可以计算。第二个可以禁用自动梯度，同时效果沿用至所有子节点。

1. import torch
2. from torch.autograd import Variable
4. # requires grad
5. # If there’s a single input to an operation that requires gradient,
6. # its output will also require gradient.
7. x = Variable(torch.randn(5, 5))
8. y = Variable(torch.randn(5, 5))
9. z = Variable(torch.randn(5, 5), requires_grad=True)
10. a = x + y
11. a.requires_grad  # False
12. b = a + z
13. b.requires_grad  # True
15. # Volatile differs from requires_grad in how the flag propagates.
16. # If there’s even a single volatile input to an operation,
17. # its output is also going to be volatile.
18. x = Variable(torch.randn(5, 5), requires_grad=True)
19. y = Variable(torch.randn(5, 5), volatile=True)
20. a = x + y
21. a.requires_grad  # False

训练过程

当然，PyTorch 还有一些其他卖点。例如你可以设定学习速率，让它以特定规则进行变化。或者你可以通过简单的训练标记允许/禁止批规范层和 dropout。如果你想要做的话，让 CPU 和 GPU 的随机算子不同也是可以的。

1. # scheduler example
2. from torch.optim import lr_scheduler
4. optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
5. scheduler = lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
7. for epoch in range(100):
8.    scheduler.step()
9.    train()
10.    validate()
12. # Train flag can be updated with boolean
13. # to disable dropout and batch norm learning
14. model.train(True)
15. # execute train step
16. model.train(False)
17. # run inference step
19. # CPU seed
20. torch.manual_seed(42)
21. # GPU seed
22. torch.cuda.manual_seed_all(42)

同时，你也可以添加模型信息，或存储/加载一小段代码。如果你的模型是由 OrderedDict 或基于类的模型字符串，它的表示会包含层名。

1. from collections import OrderedDict
3. import torch.nn as nn
5. model = nn.Sequential(OrderedDict([
6.    ('conv1', nn.Conv2d(1, 20, 5)),
7.    ('relu1', nn.ReLU()),
8.    ('conv2', nn.Conv2d(20, 64, 5)),
9.    ('relu2', nn.ReLU())
10. ]))
12. print(model)
14. # Sequential (
15. #   (conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
16. #   (relu1): ReLU ()
17. #   (conv2): Conv2d(20, 64, kernel_size=(5, 5), stride=(1, 1))
18. #   (relu2): ReLU ()
19. # )
21. # save/load only the model parameters(prefered solution)
22. torch.save(model.state_dict(), save_path)
23. model.load_state_dict(torch.load(save_path))
25. # save whole model
26. torch.save(model, save_path)
27. model = torch.load(save_path)

根据 PyTorch 文档，用 state_dict() 的方式存储文档更好。

记录

训练过程的记录是一个非常重要的部分。不幸的是，PyTorch 目前还没有像 Tensorboard 这样的东西。所以你只能使用普通文本记录 Python 了，你也可以试试一些第三方库：

• logger：https://github.com/oval-group/logger
• Crayon：https://github.com/torrvision/crayon
• tensorboard_logger：https://github.com/TeamHG-Memex/tensorboard_logger
• tensorboard-pytorch：https://github.com/lanpa/tensorboard-pytorch
• Visdom：https://github.com/facebookresearch/visdom

掌控数据

你可能会记得 TensorFlow 中的数据加载器，甚至想要实现它的一些功能。对于我来说，我花了四个小时来掌握其中所有管道的执行原理。

首先，我想在这里添加一些代码，但我认为上图足以解释它的基础理念了。

PyTorch 开发者不希望重新发明轮子，他们只是想要借鉴多重处理。为了构建自己的数据加载器，你可以从 torch.utils.data.Dataset 继承类，并更改一些方法：

1. import torch
2. import torchvision as tv
5. class ImagesDataset(torch.utils.data.Dataset):
6.    def __init__(self, df, transform=None,
7.                 loader=tv.datasets.folder.default_loader):
8.        self.df = df
9.        self.transform = transform
10.        self.loader = loader
12.    def __getitem__(self, index):
13.        row = self.df.iloc[index]
15.        target = row['class_']
16.        path = row['path']
17.        img = self.loader(path)
18.        if self.transform is not None:
19.            img = self.transform(img)
21.        return img, target
23.    def __len__(self):
24.        n, _ = self.df.shape
25.        return n
27. # what transformations should be done with our images
28. data_transforms = tv.transforms.Compose([
29.    tv.transforms.RandomCrop((64, 64), padding=4),
30.    tv.transforms.RandomHorizontalFlip(),
31.    tv.transforms.ToTensor(),
32. ])
34. train_df = pd.read_csv('path/to/some.csv')
35. # initialize our dataset at first
36. train_dataset = ImagesDataset(
37.    df=train_df,
38.    transform=data_transforms
39. )
41. # initialize data loader with required number of workers and other params
42. train_loader = torch.utils.data.DataLoader(train_dataset,
43.                                           batch_size=10,
44.                                           shuffle=True,
45.                                           num_workers=16)
47. # fetch the batch(call to `__getitem__` method)
48. for img, target in train_loader:
49.    pass

有两件事你需要事先知道：

• PyTorch 的图维度和 TensorFlow 的不同。前者的是 [Batch_size × channels × height × width] 的形式。但如果你没有通过预处理步骤 torchvision.transforms.ToTensor() 进行交互，则可以进行转换。在 transforms 包中还有很多有用小工具。
• 你很可能会使用固定内存的 GPU。对此，你只需要对 cuda() 调用额外的标志 async = True，并从标记为 pin_memory = True 的 DataLoader 中获取固定批次。

最终架构

现在我们了解了模型、优化器和很多其他细节。是时候来个总结了：

这里有一段用于解读的伪代码：

1. class ImagesDataset(torch.utils.data.Dataset):
2.    pass
4. class Net(nn.Module):
5.    pass
7. model = Net()
8. optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
9. scheduler = lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
10. criterion = torch.nn.MSELoss()
12. dataset = ImagesDataset(path_to_images)
13. data_loader = torch.utils.data.DataLoader(dataset, batch_size=10)
15. train = True
16. for epoch in range(epochs):
17.    if train:
18.        lr_scheduler.step()
20.    for inputs, labels in data_loader:
21.        inputs = Variable(to_gpu(inputs))
22.        labels = Variable(to_gpu(labels))
24.        outputs = model(inputs)
25.        loss = criterion(outputs, labels)
26.        if train:
27.            optimizer.zero_grad()
28.            loss.backward()
29.            optimizer.step()
31.    if not train:
32.        save_best_model(epoch_validation_accuracy)

结论

希望本文可以让你了解 PyTorch 的如下特点：

• 它可以用来代替 Numpy
• 它的原型设计非常快
• 调试和使用条件流非常简单
• 有很多方便且开箱即用的工具

PyTorch 是一个正在快速发展的框架，背靠一个富有活力的社区。现在是尝试 PyTorch 的好时机。

原文发布时间为：2017-10-01

本文作者：Illarion

本文来自合作伙伴“数据派THU”，了解相关信息可以关注“数据派THU”微信公众号