Pytorch使用文档

最新推荐文章于 2024-07-23 17:38:40 发布

jellyDing

最新推荐文章于 2024-07-23 17:38:40 发布

阅读量1.9k

点赞数

文章标签： pytorch

本文链接：https://blog.csdn.net/QueryWo/article/details/107615020

版权

工具专栏收录该内容

1 篇文章 0 订阅

订阅专栏

简介

先上官网和 github源码

官方简介

Tensor computation (like NumPy) with strong GPU acceleration
Deep neural networks built on a tape-based autograd system

本文默认大家对深度学习已有了解，所以主要聚焦于torch框架的使用。本文内容主要参考pytroch的官方文档，边学边写。如果有什么错误，欢迎大家指正。

1 torch.Tensor

重要概念

数据类型

torch.Tensor和torch.Storage，is_contiguous()

tensor和ndarray的关系

cpu tensor和gpu tensor

重要属性

dim()、size()、shape，type()，numel()

requires_grad，grad，is_leaf，grad_fn（自动求导）

element-wise和reduction

squeeze和unsqueeze

transpose、permute、view、reshape

torch.cat、torch.stack、torch.split、torch.chunk、torch.gather

clone、detach、new_tensor、copy_

2 torch.nn

Module

Conv

Pool、Activations、BatchNorm、RNN、LSTM、Dropout、Loss Function

nn与nn.functional的关系

3 torch.optim

optimizer

auto-grad、优化器和损失函数的关系

指定不同的学习率

4 训练数据的加载torch.util.data

5 torch的使用

torch.save和torch.load

model.train()、model.eval()

1 torch.Tensor

Tensor是torch的核心，本节将从tensor的基本概念，常用方法和属性进行介绍。

重要概念

数据类型

tensor一共支持七种CPU tensor和八种GPU tensor，可以根据自己的需要和网络类型进行选择，其中最常用的32位的float类型。此外GPU支持16位的float类型，可以有效减少模型的大小，但精度有所降低。

图像来源： https://www.cnblogs.com/icodeworld/p/11882263.html

torch.Tensor和torch.Storage，is_contiguous()

图像来源：https://www.cnblogs.com/wanghui-garcia/p/10623033.html

torch的核心是Tensor类，Tensor的数据由Storage管理，每一个tensor都对应着一个storage。为了降低存储占用量和内存操作次数，torch将内存数据和tensor分开，在某些操作中，只是tensor头信息的复制与修改，而不对实际存储空间进行操作，比如切片、索引、tensor.data()、2numpy、tensor.detach()等。不同的tensor可能对应着同一个storage，只是头信息不同，例如，split()方法，得到的不同tensor，只是头信息中的offset属性不同，代表着在对应storage中的偏移位置（类似于数组下标）。所以，在torch的使用中，需要注意不同的tensor是否共享数据内存。

import torch
x = torch.rand(9,1)
print(id(x.view(3,3).storage())==id(x.storage()))  #True

比如，view()方法常被用于改变tensor的尺寸，其复制结果与原tensor共享内存，只是对内存索引进行重新排列。所以view()方法要求输入是contiguous的。因此，tensor如果经过transpose或permute之后，想使用view()方法，（一般）需要先执行tensor.contiguous()。

import torch
x = torch.rand(3,3,3)
x = x.permute(1,2,0)
#print(x.view(3,9)) 报错，可以用reshape()进行代替
print(x.reshape(3,9))  #成功
print(x.contiguous().view(3,9)) #成功

tensor和ndarray的关系

如上面torch的官方介绍，tensor的使用几乎与Numpy无异。不同的是，tensor的数值计算过程支持GPU加速，且具有计算图和梯度等属性，但在数值上二者可以互转。

cpu tensor和gpu tensor

torch与numpy的区别之一就是torch为tensor提供了gpu加速计算，所以tensor可以分为CPU tensor和gpu tensor，分别存储于cpu内存和gpu显存中。我们可以通过cup_tensor.cuda()和gpu_tensor.cup()进行互转，其本质就是cpu内存和gpu显存之间的数据复制。而numpy数据只能在CPU上进行计算，所以gpu tensor是不能直接转为numpy数组，必须先转为cpu tensor，再转为numpy数组。

is_cuda，device

is_cuda和device就是查看tensor所在位置，gpu还是cpu，以及是在哪一块gpu或cpu上。

重要属性

dim()、size()、shape，type()，numel()

这几个是tensor的基本属性，分别是指tensor的维度数，形状，数据类型和元素数量。

requires_grad，grad，is_leaf，grad_fn（自动求导）

这几个属性都与torch的自动求导（auto-gard）有关，与tensorflow1.x不同的是，torch是动态建立计算图，计算图是一个有向无环图。在使用torch框架设计一个网络模型时，我们很少关注网络的反向传播过程，这是因为torch有自动求导机制，无需我们进行设定。首先需要说明的是，0.4.0之后的torch已经将tensor和variable进行合并，所以现在tensor也支持自动梯度。

requires_grad和grad

在我们自定义一个tensor时，requires_grad默认为False，可以通过requires_grad=True进行设置。在我们定义网络层的时候，torch自动生成的权重偏置等参数是parameter类型，属于Tensor的一个子类，其requires_grad默认为True。并且它们会自动地添加到模型的参数列表和Parameters()迭代器中，即在反向传播中，torch会通过optimizer自动更新parameter的值。我们也可以通过nn.Parameter()方法将tensor转为Parameter类型。

para = torch.nn.Parameter(torch.FloatTensor(3,3))

grad属性储存的是tensor的梯度。如果tensor的requires_grad为False，则在反向传播中，不保存其梯度，目的是为了节约内存或显存。

import torch
import torch.nn as nn
input = torch.zeros(4, 3, 5, 5)
#requires_grad默认为False
print(input.requires_grad) #False

#创建一个简单的模型，只包含一个1*1的卷积层
model = nn.Conv2d(3, 3, 1)

weight  = list(model.named_parameters())[0][1] #卷积层的权重
print(type(weight))  #<class 'torch.nn.parameter.Parameter'>
print(weight.requires_grad) #True，默认

loss = torch.nn.functional.l1_loss(model(input),torch.ones(4,3,5,5))
loss.backward()
#反向传播之后，grad保存着梯度向量
print(input.grad is None) #True
print(weight.grad is None) #False

对于tensor是否保存梯度，还和其在计算图中的位置有关。只有tensor为叶节点（叶子张量，is_leaf=True）时，才保存其梯度。因此，非叶节点是不保存梯度的，非叶节点的requires_grad是不能设为True的，如下所示。

a = torch.randn([2,2],requires_grad=True)
b = a+1
b.requires_grad=True  #RuntimeError: you can only change requires_grad flags of leaf variables.

但是如下这样b就属于叶节点了：

a = torch.randn([2,2])
b = a+1
b.requires_grad=True

这是因为，上面那个a属于叶节点，所以b属于中间变量，不是叶节点。而下面这个a的requires_grad属性默认为False，其不要求获得梯度，所以a不参与反向传播，所以计算图默认将b作为叶节点，可以通过设置requires_grad=True获得梯度。此外，非叶节点虽然无法通过设置requires_grad的方式保存梯度，但是可以使用retain_grad()方法进行设置，其效果与叶节点requires_grad=True相同，同样保存梯度到grad属性。

注：在很多博文中，作者将requires_grad与是否自动求导划等号。我认为requires_grad只是代表是否保存其梯度，在反向传播的过程中，所有的tensor都需要计算梯度，但是只保存requires_grad为True，且为叶节点（is_leaf=True）的tensor（所有的权重偏置都是叶节点）。

grad_fn

grad_fn指向创建该tensor的Function对象，用于反向传播的梯度计算之用。可以认为tensor为计算图的节点，function为计算图的边，共同构成了torch计算图，tensor的grad_fn就是计算图中tensor的输入边。因此，叶节点的grad_fn属性都为None，因为叶节点没有输入。

重要方法

data和detach()

tensor.data直接返回tensor的数据，且共享内存，但复制的数据requires_grad为False，不被autograd追踪，所以得到的是一个requires_grad为False的"叶节点"，即脱离了计算图，不参与参数更新，所以这是“不安全”的。这种不安全体现在：前向传播和反向传播时，tensor的值可能会不一致，导致梯度的计算出现错误。

与.data类似的是detach()方法，detach()同样是共享内存的复制操作，但是会被autograd追踪。如果改变了tensor的值，在训练中就会报错，是安全的。

数学操作

这部分先空着，详情可见官方文档-.-

fn和fn_

首先，需要了解一个概念“in-place"操作，in-place operation被称为原地操作，直接在原变量上进行计算。torch很多数学操作方法方法提供了in-place版本，fn代指某个方法，fn_就是相应的in-place方法。in-place方法与普通方法的最大区别是：经过in-place计算的tensor的grad_fn将被赋值，就是对应的方法。从上面对grad_fn的描述来看，这是个很大的问题，因为还有非叶子张量才有grad_fn属性，所以其梯度将不被保存，变成了一个单纯的中间向量，因此fn_方法不能被用于requires_grad = True 的叶子张量和求梯度阶段需要用到的张量，因为这与上文提到的tensor.data一样，是不安全的。

element-wise和reduction

element-wise操作是直接对每个元素进行计算，大部分数学操作都是element-wise的。需要注意的是那些tensor之间的element-wise操作，一般情况下都要求具有相同的shape，比如element-wise add和element-wise mul等等。但是torch存在着一种广播机制（Broadcast），比如“tensor+标量”就属于一种广播加法。

reduction操作是对tensor的某个维度进行计算，详见https://blog.csdn.net/gyt15663668337/article/details/95850189

squeeze和unsqueeze

squeeze可以去掉tensor所有长度为1的维度，也可以指定单个维度。unsqueeze可以在指定位置增加一个维度。

这里就用到了squeeze()的in-place方法unsqueeze_()，不然需要写成a = a.unsqueeze_(1)的形式。

a = torch.randn(2,2)
print(a.shape)   #torch.Size([2, 2])
a.unsqueeze_(1)
print(a.shape)  #torch.Size([2, 1, 2])
a.squeeze_()
print(a.shape)  #torch.Size([2, 2])

transpose、permute、view、reshape

transpose：交换两个维度的位置

permute：重新排列所有维度位置

import torch 
a = torch.randn(1,2,3,4)  #torch.Size([1, 2, 3, 4])
print(a.transpose(0,1).shape)  #torch.Size([2, 1, 3, 4])
print(a.permute(3,2,1,0).shape)  #torch.Size([4, 3, 2, 1])

需要注意，transpose和permute不会直接在原tensor进行更改，而是返回一个新的tensor，但是与原tensor共享基础数据，因此tensor将变成非contiguous的。但是，如果是执行a.permute(1,0,2,3)或a.transpose(0,1)，不会改变tensor的contiguous属性，因为并没有改变数据在内存中的排列方式。

import torch 
a = torch.randn(1,2,3,4)
print(id(a.transpose(0,1).storage())==id(a.storage()))  #True
print(id(a.permute(3,2,1,0).storage())==id(a.storage()))  #True
print(a.permute(3,2,1,0).is_contiguous()) #False
print(a.transpose(1,3).is_contiguous()) #False

view、reshape都可以用于改变tensor的尺寸，且都与原tensor共享基础数据。（即一个改变，另一个也会改变）

不同点：

首先，view只能通过tensor.view()调用，而reshape可以通过tensor.shape()和torch.reshape()两种方式（大部分tensor方法都具有这两种方式）。

view()要求tensor是contiguous的，如果不是，需要先调用tensor.contiguous()方法，上面已经说过。

reshape()比较特殊，经过实验发现，如果原tensor是contiguous的，则共享基础数据：

import torch 
a = torch.randn(4,3,2)
b = a.reshape(3,8)
a[0][0]=0
print(a,b)#a和b的值存在对应关系，b的值会随着a的改变而改变

但如果不是contiguous的，则不共享，即二者数据互不影响：

import torch 
###不是contiguous的
a = a.permute(2,1,0)
b = a.reshape(3,8)
a[0][0]=0
print(a,b)#b的值不会随着a的改变而改变

torch.cat、torch.stack、torch.split、torch.chunk、torch.gather

其中最常用的是torch.cat()，在构建神经网络时，经常使用concatenate，就可以使用cat()方法实现。

##torch.cat(tensors, dim=0, out=None)
import torch
a = torch.randn(2,3)
print(torch.cat((a,a,a),0).shape)#torch.Size([6, 3])
print(torch.cat((a,a,a),1).shape)#torch.Size([2, 9])

torch.stack：同样是合并，但与cat不同的是，stack构造了一个新维度，增加维度的长度就是stack输入的个数。stack也可以这样实现：先对每一个tensor使用unsqueeze()方法在相应位置增加一个维度，然后使用cat()进行连接。

torch.split()：按照块大小分割

torch.split(tensor, split_size_or_sections, dim=0)

split_size_or_sections 如果是int，是指每个tensor在该维度的长度，进行等分，不足时最后一个子tensor可以小于该值。也可以是一个列表，指定每一个tensor的长度。dim代表分割哪一个维度。

torch.chunk()：按照块数量分割，chunks指要分为几个子tensor

torch.chunk(input, chunks, dim=0)

torch.gather(input, dim, index, out=None) → Tensor

output的shape与index的shape相同
dim与index的关系：index的第dim维度的长度为1，其他维度与input相同
output的shape与input相比，第dim维长度变为1

out[i][j][k] = input[index[i][j][k]][j][k]  # if dim == 0
out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2

clone、detach、new_tensor、copy_

这几个方法都是进行tensor的复制，区别在于：

clone：返回一个与原tensor一模一样的tensor（数据，大小，类型，位置），不共享数据内存，但提供梯度的回溯（前面有提到）；

detach：返回一个一模一样的tensor，共享数据内存，但不提供梯度的回溯（脱离了计算图）；

new_tensor：返回一个全新的一模一样的tensor，等价于tensor.clone().detach()；

copy_:只是单纯的复制数据，dtype、device和其他设置与复制对象无关。

2 torch.nn

本节是构建深度网络所常用的模块，其中包含了大部分主流网络层，下面简单介绍常用的一些网络层和方法。

Module

torch削弱了层的概念，所有的神经网络层、网络模块、网络模型（torch.nn)都是继承与nn.Module，在torch中，单独的一个卷积层和一个网络模型是没有区别的。

Module是torch网络的基础类，所有神经网络是Module的子类，如果我们要网络模型，就需要继承nn.Module。

Module提供了很多函数（48个），一般情况下，我们只使用__init__()和forward()，定义一个含有两个卷积的简单网络结构：

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

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 4, 3)
        self.conv2 = nn.Conv2d(4, 8, 3)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

构造函数__init__用于定义forward中用到的参数，包括tensor，Parameter，网络层或网络模块。__init__的参数是可以自主定义的，只需在创建网络时传入相应的参数，比如：层数，stride，输入维度，输出维度等等等。forward()方法是我们定义网络结构的地方，输入参数一般固定为单一参数input（约定俗称），但类型可以自定义，比如：字典，列表，矩阵等。

Module其他方法请移步https://zhuanlan.zhihu.com/p/88712978，其中包括其他重要方法：cup()和cuda()、state_dict()和load_state_dict()，train()和eval()等，我就不搬运了。

Conv

torch的卷积层大同小异，他们都继承于_ConvNd类，可以看下它的构造函数参数：

    def __init__(self,
                 in_channels: int,           #输入维度
                 out_channels: int,           #输出维度
                 kernel_size: _size_1_t,      #卷积核尺寸
                 stride: _size_1_t,           #步长
                 padding: _size_1_t,          #边缘填充
                 dilation: _size_1_t,         #空洞卷积
                 transposed: bool,            #转置卷积
                 output_padding: _size_1_t,   #输出边缘填充
                 groups: int,                 #组卷积
                 bias: Optional[Tensor],      #偏置
                 padding_mode: str) :         #填充模式

Pool、Activations、BatchNorm、RNN、LSTM、Dropout、Loss Function

#pool
torch.nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)

#relu激活函数
torch.nn.ReLU(inplace=False)

#批归一化
torch.nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True)

#RNN/LSTM
torch.nn.RNN(input_size, hidden_size, num_layers=1, nonlinearity=tanh, bias=True, batch_first=False, dropout, bidirectional=False)

#dropout
torch.nn.Dropout(p=0.5, inplace=False)

#loss
#reduce为True，返回长度为batch_size的向量，代表每个样本的损失，此时size_average参数失效；否则返回一个标量
#size_average=True，返回所有样本损失的平均值；为False时，返回所有样本损失的和
torch.nn.MSELoss(reduce=False, size_average=False)

nn与nn.functional的关系

其实这两个是差不多的，不过一个是包装好的类，一个是可以直接调用的函数。这种设计是为了兼顾灵活性和快捷性。nn是创建一个对象，去维护参数和中间向量，不需要自己管理学习参数，更快捷。而nn.functional需要自己创建权重和偏置，但其灵活性更好。另外，其实nn在执行会调用nn.functional里面的方法，只是进行进一步的将方法和参数进行了封装。

一般情况下，我们会尽量使用nn，除非是没有学习参数的层，或者nn不好解决的问题。

3 torch.optim

`optimizer`

优化器是模型训练的关键，torch在optim中提供很多主流的优化器，包括SGD,动量，牛顿动量，AdaGrad，RMSProp，Adam，Lookahead（好像还没集成进来，但已经有torch版本，可以自己导入）。使用方法也非常简单，只需要我们把模型的参数绑定到优化器上即可。

optimizer = optim.SGD(model.parameters(), lr = 0.01, momentum = 0.9)

auto-grad、优化器和损失函数的关系

可能很多初学者很疑惑三者有什么关系，请看下面代码：

#input model
ouput = model(input)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()

从前面的介绍中，我们知道auto-gard主要和tensor的若干个参数有关：requires_grad、grad、is_leaf、grad_fn。假设一个tensor的is_leaf为True，requires_grad为True，且grad_fn为None时（这三个属性相互影响），在调用了loss.backward()之后，tensor的grad属性就储存了当前tensor的梯度，然后调用optimizer.step()，tensor的就会根据优化方法自动更新参数。

指定不同的学习率

在构造优化器时，不但可以对学习率进行全局指定，也可以为不同层分配不同的学习率，我们可以认为创建优化器的过程就是为整个网络模型的所有参数指制定优化方法。

optim.SGD([
                {'params': model.base.parameters()},
                {'params': model.classifier.parameters(), 'lr': 1e-3}
            ], lr=1e-2, momentum=0.9)

4 训练数据的加载torch.util.data

在深度学习模型的使用中，数据的加载是一个模型训练的重点。其中最重要的类是Dataset、DataLoader和Sample。Dataset是所使用数据集的抽象类，主要包括数据集的元数据，DataLoader是数据集的加载器，并为训练提供单线程或多线程的数据加载。Sample是采样器，用于采样数据的方法。三者的关系如下：首先我们需要定义Dataset，代表整个数据集，然后使用Dataset和采样策略sampler构造数据加载器dataloader，迭代加载数据。

DataLoader的定义如下：

class DataLoader(object):
    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None,
                 batch_sampler=None, num_workers=0, collate_fn=default_collate,
                 pin_memory=False, drop_last=False, timeout=0,
                 worker_init_fn=None)

从中也可以看出三者的关系：torch根据dataset提供的数据索引（indices），和Sample提供的数据采样策略，对batch数据的处理方法（collate_fn），采取单线程或多线程（num_workers）的方式，进行一个批次数据的加载与训练。其他参数的含义较为简单：

shuffle：数据集打乱，与Sampler互斥。原因：sampler是根据索引进行采样，本身就是一种采样策略（可以通过随机的方式实现shuffle），所以如果指定了sampler，则shuffler必须为false。

batch_sampler：与Sampler类似，但是一次只返回一个batch的indices（索引），需要注意的是，一旦指定了这个参数，那么batch_size,shuffle,sampler,drop_last就不能再制定了。原因：Batch_sampler是取整个batch的策略，而sampler + batch_size + shuffler + drop_last相当于一个batch_sampler，即如何从数据集中取一个batch的数据。（dataloader具有一个迭代方法，根据采样策略不断取数据）；

pin_memory：用于加速数据交换，设为true时，将数据设为锁页内存，即只存放于内存中。

drop_last：告诉如何处理数据集长度除于batch_size余下的数据。True就抛弃，否则保留。

很多人都有一个疑惑：sampler和batch_sampler、以及sampler和collate_fn之间是什么关系？

首先sampler只是采样一个样本的策略，比如：随机采样（等价于shuffle=true），顺序采样等，它加上其他参数（batch_size,shuffle,sampler,drop_last）会决定如何采样一个批次的数据，这与batch_sampler等价。其次，sampler（batch_sampler）是采样一个批次的数据，返回的是一个batch的索引，但每个batch的索引对应的数据格式未必是相同的，所以很多时候，需要使用collate_fn进行后处理（比如他每个batch的所有图像resize成相同大小）。这里涉及一个节约内存的问题：在训练时，并不是直接对dataset所有数据进行数据预处理，而是使用collate_fn对采样得到的单个batch进行预处理，可以有效节约内存。

5 torch的使用

torch.save和torch.load

模型的序列化与反序列化：
#定义一个简易模型

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

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(3, 3, 3)

    def forward(self, x):
        return self.conv(x)

保存与加载模型参数：

model = Model()
#保存 state_dict()为参数字典
torch.save(model.state_dict(), "model_path.pth")
#加载
model.load_state_dict(torch.load("model_path.pth"))

保存与加载整个模型，包括网络的定义等：

#save
torch.save(model, PATH)
#load
model = torch.load(PATH)

自定义保存数据：

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

#load
model = Model()
optimizer = Optimizer(...)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.train()、model.eval()

model.train()和model.eval()用于设置模型的模式，即训练模式和测试评估模式，主要改变Batch Normalization和Dropout的计算方式。

BN：训练时，使用本批次的方差和均值进行计算，并会累积计算全局方差和均值，用于推断。测试时：使用训练阶段积累的全局方差和均值。

dropout：训练时，按照指定比例 $\rho$ 随机失活；测试时：所有神经元都参与计算，最后将结果乘上 $\rho$ 。

6 多GPU训练与分布式训练

单机多卡：nn.DataParallel

使用nn.DataParallel可以将模型转为多GPU并行计算：

model = nn.DataParallel(model, device_ids=[0,1,...], output_device=None, dim=0)

其中，output_device是主gpu（负责参数更新的gpu），需要说明的是：并行计算是将一个batch的数据分到不同gpu上进行并行计算，最后需要将计算结果返回到主gpu，这里可能会出现负载不均衡问题，这是因为副卡返回到主gpu的输出是预测结果，然后在主gpu上进行损失的计算，这对主gpu的负担很大。所以可以将损失的计算过程放到各自gpu上，返回各自的损失list，最后在主gpu上进行loss.mean()，得到最终的损失，然后进行参数的更新。

多机多卡：distributedDataparallel

model = torch.nn.parallel.DistributedDataParallel(model)

与dataParallel的区别：

1. distributedDataparallel既可以单机多卡，也可以多机多卡；

2. 并且支持模型并行；

3. dataParallel单进程多线程，distributedDataparallel多进程；

4. distributedDataparallel的各进程（gpu）之间的前向传播和反向传播互不干涉，在各进程梯度计算完成之后，将梯度进行汇总平均，然后再由 rank=0 的进程，将其广播到所有进程。之后，各进程用该梯度来更新参数。所以dataParallel的梯度计算是在主gpu上进行的，distributedDataparallel的梯度计算是在各个gpu上分散进行的，只在主进程上进行梯度的平均。

需要说明一点：将各个进程的梯度汇总到主进程上，同样会出现通信的不均衡，所以可以使用更加高效的通信策略：Ring all-reduce，详见https://www.jianshu.com/p/8c0e7edbefb9

7 自定义扩展

torch.nn扩展

这个最为简单，实际上，我们所有的网络定义都属于对torch.nn的扩展，只是大部分方法是默认的。我们可以通过继承nn.Module进行网络层的扩展，并实现__init__()和forward()方法。

__init__()：实现网络层的参数定义；

forward()：实现前向计算。

torch.autograd扩展

如果我们定义的网络层是可微的，那么backward()可以由autograd机制自动完成。但如果不可微，则需要自己编写backward()方法，实现损失的反向传播。有很多网络层是不可微的，比如常用的池化层，其反向传播就是通过某种方式进行计算的，而不是求导。再比如，在hard attention中，argmax就是不可导的，但可以通过蒙特卡洛采样的方式进行模拟（或者使用softargmax代替）。还有某些激活函数，变分自编码器等等。

要自定义反向传播过程，需要继承torch.autograd.Function，下面是一个官方示例：

class LinearFunction(Function):

    # 必须是staticmethod
    @staticmethod
    # 第一个是ctx，第二个是input，其他是可选参数。
    # ctx在这里类似self，ctx的属性可以在backward中调用。
    def forward(ctx, input, weight, bias=None):
        ctx.save_for_backward(input, weight, bias)
        output = input.mm(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output


    @staticmethod
    def backward(ctx, grad_output): 
        input, weight, bias = ctx.saved_variables
        grad_input = grad_weight = grad_bias = None

        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(weight)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_output.t().mm(input)
        if bias is not None and ctx.needs_input_grad[2]:
            grad_bias = grad_output.sum(0).squeeze(0)

        return grad_input, grad_weight, grad_bias