编写高效的PyTorch代码技巧
一、PyTorch基础
PyTorch是数值计算方面其中一个很流行的库，同时也是机器学习研究方面最广泛使用的框架。在很多方面，它和NUmpy非常相似，但是它可以不需要代码做太多改变的情况下，在CPUs、GPUs、TPUs上实现计算，以及非常越容易实现分布式计算的操作。PyTorch的其中一个最重要的特征就是自动微分，它可以让需要采用梯度下降算法进行训练的机器学习算法的实现更加方便，可以更高效的自动计算函数的梯度。本文的目标是提供更好的PyTorch介绍以及讨论使用PyTorch的一些最佳实践。
1.1、张量（Tensor）：
张量就是多维数组，它和numpy的数组非常相似，但多了一些函数功能。一个张量可以存储一个标量数值、一个数值、一个矩阵......
import torch # 标量数值 a = torch.tensor(3) print(a) # tensor(3) # 数组 b = torch.tensor([1, 2]) print(b) # tensor([1, 2]) # 矩阵 c = torch.zeros([2, 2]) print(c) # tensor([[0., 0.], [0., 0.]]) # 任意维度的张量 d = torch.rand([2, 2, 2]) print(d) """ tensor([[[0.1389, 0.5866], [0.7250, 0.4697]], [[0.2261, 0.2401], [0.6850, 0.3691]]]) """
Tensor还可以高效的执行代数的运算。机器学习应用中最常见的运算就是矩阵乘法。例如希望两个随机矩阵进行相乘，维度分别是3 x 5 和 5 x 4：
import torch x = torch.randn([3, 5]) y = torch.randn([5, 4]) #矩阵相乘 a = x @ y #向量相加 z = torch.randn([3, 5]) b = x + z #将张量转换为 numpy 数组，可以调用 numpy() 方法 c = b.numpy() print(c) #numpy 数组转换为张量 print(torch.tensor(c))
1.2、自动微分：
PyTorch中比较numpy最大的有点就是可以实现自动微分，这对于优化神经网络参数的应用非常有帮助。
假设现在有一个复合函数：g(u(x))，为了计算 g 对 x 的导数，这里可以采用链式法则，即
而PyTorch可以自动实现这个求导的过程。为了在PyTorch中计算导数，首先要创建一个张量（Tensor），并设置requires_grad = True，然后利用张量运算来定义函数，这里设 u 是一个二次方的函数，而 g 是一个简单的线性函数，代码如下所示：
import torch x = torch.tensor(1.0, requires_grad=True) #x的初始值为1 def u(x): return x * x #u(x) = x^2 def g(u): return -u #g(u) = -u = -(x * x) g'(u) = -2*x #在 PyTorch 中调用梯度函数 dgdx = torch.autograd.grad(g(u(x)), x)[0] print(dgdx) # tensor(-2.)
import torch x = torch.tensor(1.0, requires_grad=True) # x的初始值为1 def u(x1): return x1 * x1 # u(x) = x^2 def g(u1): return -u1 # g(u) = -u = -(x * x) g'(u) = -2*x uu = u(x) gg = g(uu) gg.backward() print(x.grad) #tensor(-2.)
1.3、拟合曲线
为了展现自动微分有多么强大，这里介绍另一个例子。首先假设我们有一些服从一个曲线（也就是 f(x) = 5x^2 + 3）的样本，然后希望基于这些样本来评估这个函数 f(x) 。我们先定义一个带参数的函数：
函数的输入是 x ，然后 w 是参数，目标是找到合适的参数，使得下列式子成立：
实现的一个方法可以是通过优化下面的损失函数来说实现：
尽管这个问题里有一个正式的函数（即 f(x) 是一个具体的函数），但这里我们还是采用一个更加通用的方法，可以应用到任何一个可微分的函数，并采用随机梯度下降法，即通过计算 L(w) 对于每个参数 w 的梯度的平均值，然后不断从相反反向移动。
利用PyTorch实现的代码如下所示：
import torch # 假设我们知道期望的函数是一个二阶多项式，我们分配一个大小为3的向量来保持系数，并用随机噪声初始化它。 w = torch.tensor(torch.randn([3, 1]), requires_grad=True) # 我们使用Adam优化器，学习率设置为0.1，以最小化损失。 opt = torch.optim.Adam([w], 0.1) def model(x): # 我们把yhat定义为y的估计值。 f = torch.stack([x * x, x, torch.ones_like(x)], 1) # torch.ones_like(x):返回一个填充了标量值1的张量，其大小与x相同 # print(f.shape) # (100,3) yhat = torch.squeeze(f @ w, 1) # print(yhat.shape) return yhat def compute_loss(y, yhat): # 损失的定义是y的估计值和它的真实值之间的误差均值的平方距离。 loss = torch.nn.functional.mse_loss(yhat, y) return loss def generate_data(): # 根据真函数生成一些训练数据 x = torch.rand(100) * 20 - 10 # (1,100) y = 5 * x * x + 3 return x, y def train_step(): x, y = generate_data() yhat = model(x) # 正向传播 loss = compute_loss(y, yhat) opt.zero_grad() loss.backward() opt.step() # step()方法，这个方法会更新所有的参数 for _ in range(1000): train_step() print(w.detach().numpy()) # [4.9924135, 0.00040895029, 3.4504161]
上述只是 PyTorch 可以做的事情的冰山一角。很多问题，比如优化一个带有上百万参数的神经网络，都可以用 PyTorch 高效的用几行代码实现，PyTorch 可以跨多个设备和线程进行拓展，并且支持多个平台。
二、将模型封装为模块
在之前的例子中，我们构建模型的方式都是直接实现张量间的运算操作。但是为了让代码看起来更加有组织，推荐采用PyTorch的 Modules 模块。一个模块实际上就是一个包含参数和压缩模型运算的容器。
比如，如果想实现一个线性模型 y = ax + b ，那么实现的代码可以如下所示：
import torch class Net(torch.nn.Module): def __init__(self): super().__init__() self.a = torch.nn.Parameter(torch.rand(1)) self.b = torch.nn.Parameter(torch.rand(1)) def forward(self, x): yhat = self.a * x + self.b return yhat
使用的例子如下所示，需要实例化声明的模型，并且像调用函数一样使用它：
x = torch.arange(100, dtype=torch.float32) net = Net() y = net(x)
参数都是设置 requires_grad 为 True 的张量。通过模型的 parameters() 方法可以很方便的访问和使用参数，如下所示：
for p in net.parameters(): print(p)
例子：
现在，假设是一个未知的函数（y = 5x + 3 + n），注意这里的 n 是表示噪声，然后希望优化模型参数来拟合这个函数，首先可以简单从这个函数进行采样，得到一些样本数据：
x = torch.arange(100, dtype=torch.float32) / 100
y = 5 * x + 3 + torch.rand(100) * 0.3
和上一个例子类似，需要定义一损失函数并优化模型的参数，如下所示：
criterion = torch.nn.MSELoss() optimizer = torch.optim.SGD(net.parameters(), lr=0.01) for i in range(10000): net.zero_grad() yhat = net(x) loss = criterion(yhat, y) loss.backward() optimizer.step() print(net.a, net.b) # Should be close to 5 and 3
汇总：
1 class Net(nn.Module): 2 def __init__(self): 3 super(Net, self).__init__() 4 self.a = torch.nn.Parameter(torch.rand(1)) 5 self.b = torch.nn.Parameter(torch.rand(1)) 6 7 def forward(self, x): 8 yhat = self.a * x + self.b 9 return yhat
1 import torch 2 from wupiao1 import Net 3 4 5 x = torch.arange(100, dtype=torch.float32) / 100 6 y = 5 * x + 3 + torch.rand(100) * 0.3 7 8 net = Net() 9 10 criterion = torch.nn.MSELoss() 11 optimizer = torch.optim.SGD(net.parameters(), lr=0.01) 12 13 for i in range(10000): 14 net.zero_grad() 15 yhat = net(x) 16 loss = criterion(yhat, y) 17 loss.backward() 18 optimizer.step() 19 20 print(net.a, net.b) # Should be close to 5 and 3
重定义模块：
在PyTorch中已经实现了很多预定义好的模块。比如 torch.nn.Linear 就是一个类似上述例子中定义的一个更加通用的线性函数，所以我们可以采用这个函数来重写我们的模块代码，如下所示：
class Net(torch.nn.Module): def __init__(self): super().__init__() self.linear = torch.nn.Linear(1, 1) def forward(self, x): yhat = self.linear(x.unsqueeze(1)).squeeze(1) return yhat
这里用到了两个函数， squeeze 和 unsqueeze ，主要是 torch.nn.Linear 会对一批向量而不是数值进行操作。
同样，默认调用 parameters() 会返回其所有子模块的参数：
net = Net() for p in net.parameters(): print(p)
汇总：
1 class Net(torch.nn.Module): 2 def __init__(self): 3 super(Net, self).__init__() 4 self.linear = torch.nn.Linear(1, 1) 5 6 def forward(self, x): 7 # print(x.shape) #torch.Size([100]) 标量 8 # print((x.unsqueeze(1)).shape) #torch.Size([100, 1]) 向量 9 b = self.linear(x.unsqueeze(1)) 10 yhat = b.squeeze(1) 11 # print(yhat.shape) #torch.Size([100]) 标量 12 return yhat
1 import torch 2 from wupiao1 import Net 3 4 5 x = torch.arange(100, dtype=torch.float32) / 100 6 y = 5 * x + 3 + torch.rand(100) * 0.3 7 8 net = Net() 9 10 criterion = torch.nn.MSELoss() 11 optimizer = torch.optim.SGD(net.parameters(), lr=0.01) 12 13 for i in range(1): 14 net.zero_grad() 15 yhat = net(x) 16 loss = criterion(yhat, y) 17 loss.backward() 18 optimizer.step() 19 20 for p in net.parameters(): 21 print(p)
当然也有一些预定义的模块（torch.nn.Sequential）是作为包容其他模块的容器，最常见的就是 torch.nn.Sequential ，它的名字就暗示了它主要用于堆叠对个模块（或者网络层），例如，堆叠两个线性网络层，中间是一个非线性函数 RuLU ，如下所示：
model = torch.nn.Sequential( torch.nn.Linear(64, 32), torch.nn.ReLU(), torch.nn.Linear(32, 10), )
torch.nn.Sequential的详细使用见：nn.Sequential类-使用Sequential类来自定义顺序连接模型
三、广播机制的优缺点
优点：
PyTorch支持广播的元素积运算。正常情况下，当想执行类似加法和乘法操作的时候，你需要确认操作数的形状是匹配的，比如无法进行一个 [3, 2] 大小的张量和 [3, 4] 大小的张量的加法操作。
但是存在一种特殊的情况：只有单一维度的时候，PyTorch 会隐式的根据另一个操作数的维度来拓展只有单一维度的操作数张量。因此，实现 [3,2] 大小的张量和 [3,1] 大小的张量相加的操作是合法的。如下所示：
import torch a = torch.tensor([[1., 2.], [3., 4.]]) b = torch.tensor([[1.], [2.]]) # c = a + b.repeat([1, 2]) """ tensor([[1., 1.], [2., 2.]]) """ c = a + b print(c) """ tensor([[2., 3.], [5., 6.]]) """
广播机制可以实现隐式的维度复制操作（repeat 操作），并且代码更短，内存使用上也更加高效，因为不需要存储复制的数据的结果。这个机制非常适合用于结合多个维度不同的特征的时候。
为了拼接不同维度的特征，通常的做法是先对输入张量进行维度上的复制，然后拼接后使用非线性激活函数。整个过程的代码实现如下所示：
缺点：
四、使用好重载的运算符
五、采用TorchScript优化运行时间
六、构建高效的自定义数据加载类
七、PyTorch的数值稳定性
参考：
参考1：编写高效的PyTorch代码技巧
浙公网安备 33010602011771号