PyTorch爬坑指南

发表于 更新于

用PyTorch也有个把年头了，但是偶尔还是会出现各种各样的bug，当bug积累到一定数量就会引起质变——写成一个笔记来长记性……

Tensor的转置

tensor的转置有很多方式，比如.transpose、.transpose_是最常用的（带下划线的是在原tensor上做改变），但是这个转置方式有个限制就是仅能对两个维度进行操作，是个“根正苗红”的转置函数。

有时候需要将更多的维度进行交换，这个时候就可以用.permuate函数，该函数接受一个维度列表，然后根据列表中的维度的排列顺序更新tensor的维度排列。

nn.Linear的坑人之处

全连接层常用的nn.Linear实际上就是将weight和输入的inputs做一个矩阵乘法，然后加上bias（如果有的话）。所以当我们想在模型中添加一个纯矩阵运算的时候，往往会想到用不带bias的nn.Linear。然而谁知道这最简单的模块也会存在坑……

假设我们有一个输入形状为batch×32的输入矩阵，然后想通过矩阵运算将其变成形状为batch×16的输出，那么就会就会这么做：

1
2
3
4
inputs = torch.rand(batch, 32)
matrix = nn.Linear(32, 16, bias=False)
outputs = matrix(inputs)
# outputs.shape == torch.Size([batch, 16])

我们一直这么使用，这样的实现方式也是十分自然的。可是当我们现在手上有一个具有特殊性质的张量pretrained，并希望用这个张量填充matrix，我们会想当然地这么做：

1
2
3
4
# 反例！
# pretrained.shape == torch.Size([32, 16])
matrix.weight.data = pretrained
outputs = matrix(inputs)

不出意外，PyTorch会抛出一个RuntimeError，并说运算过程中两个做乘法的矩阵不满足矩阵乘法的要求。

经过仔细的debug，我终于发现了问题的所在。让我们看一下在PyTorch源码中，nn.Linear是怎么初始化的：

1
2
3
4
5
6
7
8
9
10
11
12
class Linear(Module):
    ...
    def __init__(self, in_features, out_features, bias=True):
        super(Linear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.Tensor(out_features, in_features))
        if bias:
            self.bias = Parameter(torch.Tensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

可以发现，self.weight的形状是我们之前预期形状的转置！╮(╯▽╰)╭所以正确的填充方法应该是这样的：

1
2
3
4
# 正例√
# pretrained.shape == torch.Size([32, 16])
matrix.weight.data = pretrained.transpose(0, 1)
outputs = matrix(inputs)

升维和降维

升维.unsqueeze是在指定的维度上插入一维，但是数据实际上没有发生变化。比如下面得示例：

1
2
3
4
5
6
7
8
9
10
11
12
In [1]: a
Out[1]:
tensor([[[-0.7908, -1.4576, -0.3251],
         [-1.2053,  0.3667,  0.9423],
         [ 0.0517,  0.6051, -0.1360],
         [ 0.8666, -1.4679, -0.4511]]])

In [2]: a.shape
Out[2]: torch.Size([1, 4, 3])

In [3]: a.unsqueeze(2).shape
Out[3]: torch.Size([1, 4, 1, 3])

降维.squeeze是升维得逆操作，会消除所有维度为1的维度：

1
2
3
4
5
In [1]: a.shape
Out[1]: torch.Size([1, 4, 3])

In [2]: a.squeeze().shape
Out[2]: torch.Size([4, 3])

Tensor的contiguous

一个Tensor执行转置操作后，实际上数据并没有发生任何变化，只是读取数据的顺序变化了。这是一种节约空间和运算的方法。

不过转置之后的数据如果需要进行.view等操作的话，由于数据在逻辑上并不连续，因此需要手动调用contiguous让数据恢复连续存储的模样。

tensor.expand函数

PyTorch有和Numpy一样的广播方式，可以将一个较小维度的数据广播为高维数据，其中有一个叫expand的函数就能实现这个功能。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
In [1]: a = torch.arange(0, 4).view(1, -1)

In [2]: a
Out[2]: tensor([[0, 1, 2, 3]])

In [3]: a.expand(5, -1)
Out[3]: 
tensor([[0, 1, 2, 3],
        [0, 1, 2, 3],
        [0, 1, 2, 3],
        [0, 1, 2, 3],
        [0, 1, 2, 3]])

In [4]: b = torch.arange(0, 4).view(-1, 1)

In [5]: b
Out[5]: 
tensor([[0], [1], [2], [3]])

In [6]: b.expand(-1, 5)
Out[6]: 
tensor([[0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1],
        [2, 2, 2, 2, 2],
        [3, 3, 3, 3, 3]])

In [7]: c = torch.arange(0, 4).view(2, 2)

In [8]: c
Out[8]: 
tensor([[0, 1],
        [2, 3]])

In [9]: c.expand(2, -1, -1)
Out[9]: 
tensor([[[0, 1],
         [2, 3]],

        [[0, 1],
         [2, 3]]])

只不过这个函数存在一些需要注意的点。

  1. -1表示原来的数据宽度
  2. 如果需要在某一个已存在的维度上进行扩展，那么该维度的tensor宽度只能是1。换句话说，数据宽度为1的tensor可以扩展为任意维度。
  3. 如果需要将原来的tensor扩展到更高维，那么新拓展的维度一定需要放在前面（高维），如果需要改变排序方式，可以用类似.permute的其他函数。
  4. __扩展之后的数据只是一种“视图”，而不分配新的空间__。如果原来的tensor改变，那么扩展后的tensoer也会发生变化。就像下面这个例子。所以，必要时可以用.copy。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
In [10]: a
Out[10]: tensor([[0, 1, 2, 3]])

In [11]: a_view = a.expand(4, -1)

In [12]: a_view
Out[12]: 
tensor([[0, 1, 2, 3],
        [0, 1, 2, 3],
        [0, 1, 2, 3],
        [0, 1, 2, 3]])

In [13]: a *= 10

In [14]: a_view
Out[14]: 
tensor([[ 0, 10, 20, 30],
        [ 0, 10, 20, 30],
        [ 0, 10, 20, 30],
        [ 0, 10, 20, 30]])

Loss Functions的参数格式到底是啥啊

这个是早期经常会遇到的情况，咋一会要FloatTensor，一会就要LongTensor了？要解决这个问题，还是需要深入了解不同的loss function到底在做啥。

BCELoss

二分类的交叉熵，计算公式如下：

$$ \sum_{i}{-y_ilog\hat{y_i}-(1-y_i)log(1-\hat{y_i})} $$

这样就不难看出 $ y_i $ 和 $ \hat{y_i} $ 都需要是FloatTensor了。如果是不同的tensor进行运算会出错：

1
Expected object of type torch.FloatTensor but found type torch.LongTensor

CrossEntropyLoss

上面的BCELoss是计算二分类情况下的损失函数，那么当遇到多分类问题的时候，也想用交叉熵来计算的话，就需要用到CrossEntropyLoss了。

不论是二分类还是多分类，都可以用最简单的L2损失来计算，但是如果使用Sigmoid作为激活函数的话容易导致梯度消失，导致效果不好

CrossEntropyLoss输入参数的形状一直让我很头疼。比如最简单的例子：

1
2
3
4
5
6
In [1]: inputs = torch.randn(3, 5)

In [2]: targets = torch.randint(0, 5, (3, ))

In [3]: nn.CrossEntropyLoss()(inputs, targets)
Out[3]: tensor(1.4814)

一切安好。这很容易让我们想到，只要batch那一维度对应上就没有太大的问题了。于是对应到高维操作，假设有一个任务，输入张量inputs的形状是(seq_len, batch, dim)，目标张量targets的形状是(seq_len, batch)，尝试一下：

1
2
3
4
5
6
7
8
9
10
11
12
13
In [4]: seq_len, batch, dim = 3, 4, 5

In [5]: inputs = torch.randn(seq_len, batch, dim)

In [6]: targets = torch.randint(0, dim, (seq_len, batch))

In [7]: nn.CrossEntropyLoss()(inputs, targets)
---------------------------------------------------------------------------
ValueError  Traceback (most recent call last)
<ipython-input-46-37b720a7a2f6> in <module>()
----> 1 nn.CrossEntropyLoss()(inputs, targets)
    ...
ValueError: Expected target size (3, 5), got torch.Size([3, 4])

居然报错了……

??

实际上之前偷懒的解决方式都是用view把inputs碾平到2维……暂且不考虑CrossEntropyLoss的内部操作，碾平的操作本身就很不优雅。于是参考高阶处理方式。

官方文档给出的API说明是这样的：输入参数有Input和Target，前者的形状可以是$(minibatch,C,d_1,d_2,\dots,d_k)$，后者的形状是$(minibatch,d_1,d_2,\dots,d_k)$，其中$C$表示需要模型分类的类别数。

根据API的定义，之前的任务应该这样实现：

1
2
3
4
# inputs = torch.randn(seq_len, batch, dim)
# targets = torch.randint(0, dim, (seq_len, batch))
In [8]: nn.CrossEntropyLoss()(inputs.permute(1, 2, 0), targets.permute(1, 0))
Out[8]: tensor(2.0654)

自定义网络一定要初始化权值

权值不初始化会出事的……

有时候会遇到一些结构清奇的网络，各种参数纠缠在一起，因此torch自带的函数可能不够用，那就需要自定义网络结构了。一般会从最底层的矩阵开始构建，就像tensorflow一样。比如定义一个不带偏置项的全连接层可以用nn.Linear(xxx, xxx, bias=False)。但特殊情况下甚至连全连接层内部的细节都要自己定义，这个时候就需要参考torch的源码了，比如全连接层的部分源码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Linear(Module):
    def __init__(self, in_features, out_features, bias=True):
        super(Linear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.Tensor(out_features, in_features))
        if bias:
            self.bias = Parameter(torch.Tensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, input):
        return F.linear(input, self.weight, self.bias)

在这个例子中已经能学到很多了，这里重点参考Parameters的使用（out_feature在前）和权值初始化函数reset_parameters的内部细节。

之前一次机器学习的作业，相同的模型，PyTorch版始终达不到Tensorflow版的准确率，大概就是权值初始化两者存在差异吧。

CUDA的一些坑

自定义的网络不会自动转cuda？？

之前有用过自己先定义了一个简单的底层网络，方便上层网络的构建。这样定义出来的网络在CPU上运行的好好的，但是当我把模型迁移到GPU上时却发现模型出现了类型不符的错误：

1
Expected object of type torch.FloatTensor but found type torch.cuda.FloatTensor for argument

后来在调试的时候发现对上层网络调用.cuda后，下层网络却还是处于device(type='cpu')！虽然解决这个问题最简单的方法就是把底层网络的定义放到上层网络中，但是这样不利于代码的维护，并且在不同网络中共享模块也会变得很复杂。

PyTorch在CPU和GPU数据上的迁移的确没有tensorflow来的方便……

在GitHub上看到很多类似的问题的解决方法，其中有一个点子很不错：

1
2
3
4
if torch.cuda.available():
    import torch.cuda as t
else:
    import torch as t

如果一开始写代码的时候就考虑到自动CUDA的话，这的确是个好办法，如果已经写了臃肿的代码再改写还是会很繁琐。

最终我选择的方案还是重写.cuda函数：

1
2
3
def cuda(self, device=None):
    r"""Moves all model parameters and buffers to the GPU..."""
    return self._apply(lambda t: t.cuda(device))

在.cuda函数中对子模块也调用.cuda就可以了。

CUDA数据不能直接转numpy数据

如果一个数据已经迁移到GPU上了，那么如果要将运算结果转为符合numpy特性的数据就不是这么直接的了，一般会报错：

1
TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

解决方案报错信息已经给了。当然还可以将之后需要numpy来帮忙的运算改成PyTorch支持的运算，这样速度还更快呢，毕竟条条大路通罗马。

爱因斯坦简记法

偶然间看到一个十分有趣的矩阵、向量、标量的计算方法：torch.einsum。这是由爱因斯坦提出的简记法。

假设有矩阵：$A_{i \times j}$、$B_{i \times j}$、$C_{j \times k}$。

矩阵乘法可以表示为torch.einsum("ij, jk -> ik", A, C)；

逐点乘法可以表示为torch.einsum("ij, ij -> ij", A, B)；

矩阵转置可表示为torch.einsum("ij -> ji", A)；

如果矩阵是方阵的话，提取对角线元素可以用torch.einsum('ii -> i', sq)；获得方阵的迹可以用torch.einsum('ii ->', sq)；

矩阵降维可以用：torch.einsum('ij -> i', sq)或者torch.einsum('ij -> j', sq)；更高维也是如此；

其实Numpy也有这样的函数：np.einsum，并且似乎比PyTorch的计算速度快很多。下面是测试代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
In [1]: %timeit _ = torch.einsum('ij, jk -> ik', aten, bten)
70.5 µs ± 2.4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [2]: %timeit _ = torch.einsum('ij, jk -> ik', a, b)
66.7 µs ± 2.04 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [3]: %timeit _ = np.einsum('ij, jk -> ik', a, b)
11.3 µs ± 501 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [4]: %timeit _ = torch.matmul(a, b)
6.58 µs ± 388 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [5]: %timeit _ = np.matmul(a, b)
22.6 µs ± 1.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

虽然使用einsum十分简洁，但是从速度上考虑，还是老老实实用最初级的API吧……