PyTorch Code Snippets

1. 卷积计算

代码如下：

import torch
filter = torch.nn.Conv2d(2, 1, 2, 1, 1, 3, bias=False)

filter.weight[0][0][0][0]=2
filter.weight[0][0][0][1]=1
filter.weight[0][0][1][0]=1
filter.weight[0][0][1][1]=-1

filter.weight[0][1][0][0]=-3
filter.weight[0][1][0][1]=2
filter.weight[0][1][1][0]=-2
filter.weight[0][1][1][1]=3

x = torch.FloatTensor([[[[4,-3,2,2],[3,5,-2,3],[0,-4,1,7],[-8,2,9,2]], [[1,-5,3,7],[3,6,9,-2],[2,4,2,-1],[2,-2,7,0]]]])

filter(x)

2. 梯度计算

2.1 requires_grad

如果一个Tensor的requires_grad为True，则需要计算梯度，为False，则不用计算。

2.2 with torch.no_grad()

torch.no_grad是一个用于禁止梯度计算的上下文管理器，只作用于当前线程。

在Inference阶段，不会调用Tensor.backward()，因而无需计算梯度。

禁止梯度计算能够节省很多内存和显存。

>>> x = torch.tensor([1], requires_grad=True)
>>> with torch.no_grad():
...   y = x * 2
>>> y.requires_grad
False

2.3 @torch.no_grad()

也可以作为修饰器来使用，如下：

>>> @torch.no_grad()
... def doubler(x):
...     return x * 2
>>> z = doubler(x)
>>> z.requires_grad
False