PyTorch常用代码段

记录一些常用的Pytorch代码,方便后续查找。

编码风格

我们试图遵循 Google 的 Python 编程风格。请参阅 Google 提供的优秀的 python 编码风格指南(见参考资料)。

在这里,我们会给出一个最常用命名规范小结:

TypeConventionExample
Packages & Moduleslower_with_underfrom prefetch_generator import BackgroundGenerator
ClassesCapWordsclass DataLoader
ConstantsCAPS_WITH_UNDERBATCH_SIZE=16
Instanceslower_with_underdataset = Dataset
Methods & Functionslower_with_under()def visualize_tensor()
Variableslower_with_underbackground_color=’Blue’

常用模块

常用的模块有:

NameDescriptionUsed for
torchBase Framework for working with neural networkscreating tensors, networks and training them using backprop
torchvisiontododata preprocessing, augmentation, postprocessing
Pillow (PIL)Python Imaging LibraryLoading images and storing them
NumpyPackage for scientific computing with PythonData preprocessing & postprocessing
prefetch_generatorLibrary for background processingLoading next batch in background during computation
tqdmProgress barProgress during training of each epoch
torchsummaryKeras summary for PyTorchDisplays network, it’s parameters and sizes at each layer
tensorboardxTensorboard without tensorflowLogging experiments and showing them in tensorboard

文件组织

不要将所有的层和模型放在同一个文件中。最好的做法是将最终的网络分离到独立的文件(networks.py)中,并将层、损失函数以及各种操作保存在各自的文件中(layers.py,losses.py,ops.py)。最终得到的模型(由一个或多个网络组成)应该用该模型的名称命名(例如,yolov3.py,DCGAN.py),且引用各个模块。

主程序、单独的训练和测试脚本应该只需要导入带有模型名字的 Python 文件。

基础配置

检查 PyTorch 版本

1
2
3
4
5
torch.__version__               # PyTorch version
torch.version.cuda # Corresponding CUDA version
torch.backends.cudnn.version() # Corresponding cuDNN version
torch.cuda.get_device_name(0) # GPU type
torch.cuda.is_available() # CUDA

固定随机种子

1
2
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

指定程序运行在特定 GPU 卡上

在命令行指定环境变量

1
CUDA_VISIBLE_DEVICES=0,1 python train.py

或在代码中指定

1
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

设置为 cuDNN benchmark 模式

Benchmark 模式会提升计算速度,但是由于计算中有随机性,每次网络前馈结果略有差异。

1
torch.backends.cudnn.benchmark = True

如果想要避免这种结果波动,设置

1
torch.backends.cudnn.deterministic = True

清除 GPU 存储

有时 Control-C 中止运行后 GPU 存储没有及时释放,需要手动清空。在 PyTorch 内部可以

1
torch.cuda.empty_cache()

或在命令行可以先使用 ps 找到程序的 PID,再使用 kill 结束该进程

1
ps aux | grep pythonkill -9 [pid]

或者直接重置没有被清空的 GPU

1
nvidia-smi --gpu-reset -i [gpu_id]

张量处理

张量基本信息

1
2
3
tensor.type()   # Data type
tensor.size() # Shape of the tensor. It is a subclass of Python tuple
tensor.dim() # Number of dimensions.

数据类型转换

1
2
3
4
5
6
7
8
# Set default tensor type. Float in PyTorch is much faster than double.
torch.set_default_tensor_type(torch.FloatTensor)

# Type convertions.
tensor = tensor.cuda()
tensor = tensor.cpu()
tensor = tensor.float()
tensor = tensor.long()

torch.Tensor 与 np.ndarray 转换

1
2
3
4
5
6
# torch.Tensor -> np.ndarray.
ndarray = tensor.cpu().numpy()

# np.ndarray -> torch.Tensor.
tensor = torch.from_numpy(ndarray).float()
tensor = torch.from_numpy(ndarray.copy()).float() # If ndarray has negative stride

torch.Tensor 与 PIL.Image 转换

PyTorch 中的张量默认采用 N×D×H×W 的顺序,并且数据范围在 [0, 1],需要进行转置和规范化。

1
2
3
4
5
6
7
# torch.Tensor -> PIL.Image.
image = PIL.Image.fromarray(torch.clamp(tensor * 255, min=0, max=255).byte().permute(1, 2, 0).cpu().numpy())
image = torchvision.transforms.functional.to_pil_image(tensor) # Equivalently way

# PIL.Image -> torch.Tensor.
tensor = torch.from_numpy(np.asarray(PIL.Image.open(path))).permute(2, 0, 1).float() / 255
tensor = torchvision.transforms.functional.to_tensor(PIL.Image.open(path)) # Equivalently way

np.ndarray 与 PIL.Image 转换

1
2
3
4
5
# np.ndarray -> PIL.Image.
image = PIL.Image.fromarray(ndarray.astypde(np.uint8))

# PIL.Image -> np.ndarray.
ndarray = np.asarray(PIL.Image.open(path))

从只包含一个元素的张量中提取值

这在训练时统计 loss 的变化过程中特别有用。否则这将累积计算图,使 GPU 存储占用量越来越大。

1
value = tensor.item()

张量形变

张量形变常常需要用于将卷积层特征输入全连接层的情形。相比 torch.view,torch.reshape 可以自动处理输入张量不连续的情况。

1
tensor = torch.reshape(tensor, shape)

view

返回一个有相同数据但是不同形状的新的向量。

1
torch.Tensor.view(*args) → Tensor

返回的装两必须与原张量有相同的数据和相同的元素个数,但是可以有不同的尺寸。

参数:

  • args (torch.Size or int…) - 理想的指定尺寸
1
2
3
4
5
6
x = torch.randn(4, 4)
x.size()
torch.Size([4, 4])
y = x.view(16)
y.size()
torch.Size([16])

打乱顺序

1
tensor = tensor[torch.randperm(tensor.size(0))]  # Shuffle the first dimension

水平翻转

PyTorch 不支持 tensor[::-1] 这样的负步长操作,水平翻转可以用张量索引实现。

1
# Assume tensor has shape N*D*H*W.tensor = tensor[:, :, :, torch.arange(tensor.size(3) - 1, -1, -1).long()]

复制张量

有三种复制的方式,对应不同的需求。

1
2
3
4
# Operation                 |  New/Shared memory | Still in computation graph |
tensor.clone() # | New | Yes |
tensor.detach() # | Shared | No |
tensor.detach.clone()() # | New | No |

拼接张量

注意 torch.cat 和 torch.stack 的区别在于 torch.cat 沿着给定的维度拼接,而 torch.stack 会新增一维。例如当参数是 3 个 10×5 的张量,torch.cat 的结果是 30×5 的张量,而 torch.stack 的结果是 3×10×5 的张量。

1
2
tensor = torch.cat(list_of_tensors, dim=0)
tensor = torch.stack(list_of_tensors, dim=0)

压缩张量

1
torch.squeeze(input, dim=None, out=None) → Tensor

除去输入张量input中数值为1的维度,并返回新的张量。如果输入张量的形状为(A×1×B×C×1×D),那么输出张量的形状为(A×B×C×D)。

当通过dim参数指定维度时,维度压缩操作只会在指定的维度上进行。如果输入向量的形状为(A×1×B),squeeze(input, 0)会保持张量的维度不变,只有在执行squeeze(input, 1)时,输入张量的形状会被压缩至(A×B)。

如果一个张量只有1个维度,那么它不会受到上述方法的影响。

输出的张量与原张量共享内存,如果改变其中的一个,另一个也会改变

参数:

  • input (Tensor) – 输入张量
  • dim (int, optional) – 如果给定,则只会在给定维度压缩
  • out (Tensor, optional) – 输出张量
1
2
3
4
5
6
7
8
9
10
11
12
x = torch.zeros(2, 1, 2, 1, 2)
x.size()
torch.Size([2, 1, 2, 1, 2])
y = torch.squeeze(x)
y.size()
torch.Size([2, 2, 2])
y = torch.squeeze(x, 0)
y.size()
torch.Size([2, 1, 2, 1, 2])
y = torch.squeeze(x, 1)
y.size()
torch.Size([2, 2, 1, 2])

相反的,torch.unsqueeze(n)则是在第n维增加一个维数=1,如下,表示在原始的b的第二维增加一维,则尺寸变为1 3 1 * 2

1
2
3
4
5
6
7
8
>>> b.unsqueeze(2).size()
torch.Size([1, 3, 1, 2])
>>> b.unsqueeze(2)
tensor([[[[1., 2.]],

[[3., 4.]],

[[5., 6.]]]])

置换张量维度

1
torch.Tensor.permute(*dims)

将执行本方法的张量的维度换位。

参数:

  • dim (int) - 指定换位顺序

例子:

1
2
3
4
5
x = torch.randn(2, 3, 5)
x.size()
torch.Size([2, 3, 5])
x.permute(2, 0, 1).size()
torch.Size([5, 2, 3])

将整数标记转换成独热(one-hot)编码

PyTorch 中的标记默认从 0 开始。

1
2
3
N = tensor.size(0)
one_hot = torch.zeros(N, num_classes).long()
one_hot.scatter_(dim=1, index=torch.unsqueeze(tensor, dim=1), src=torch.ones(N, num_classes).long())

得到非零/零元素

1
2
3
4
torch.nonzero(tensor)               # Index of non-zero elements
torch.nonzero(tensor == 0) # Index of zero elements
torch.nonzero(tensor).size(0) # Number of non-zero elements
torch.nonzero(tensor == 0).size(0) # Number of zero elements

张量扩展

1
2
# Expand tensor of shape 64*512 to shape 64*512*7*7.
torch.reshape(tensor, (64, 512, 1, 1)).expand(64, 512, 7, 7)

矩阵乘法

1
2
3
4
5
6
7
8
# Matrix multiplication: (m*n) * (n*p) -> (m*p).
result = torch.mm(tensor1, tensor2)

# Batch matrix multiplication: (b*m*n) * (b*n*p) -> (b*m*p).
result = torch.bmm(tensor1, tensor2)

# Element-wise multiplication.
result = tensor1 * tensor2

计算两组数据之间的两两欧式距离

1
2
3
4
5
6
# X1 is of shape m*d.
X1 = torch.unsqueeze(X1, dim=1).expand(m, n, d)
# X2 is of shape n*d.
X2 = torch.unsqueeze(X2, dim=0).expand(m, n, d)
# dist is of shape m*n, where dist[i][j] = sqrt(|X1[i, :] - X[j, :]|^2)
dist = torch.sqrt(torch.sum((X1 - X2) ** 2, dim=2))

模型定义

我们建议将网络分解为更小的可复用的片段。一个 nn.Module 网络包含各种操作或其它构建模块。损失函数也是包含在 nn.Module 内,因此它们可以被直接整合到网络中。

继承 nn.Module 的类必须拥有一个forward方法,它实现了各个层或操作的前向传导。

一个 nn.module 可以通过self.net(input)处理输入数据。在这里直接使用了对象的call()方法将输入数据传递给模块。

1
output = self.net(input)

卷积层

最常用的卷积层配置是

1
conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True)

如果卷积层配置比较复杂,不方便计算输出大小时,可以利用如下可视化工具辅助

链接:https://ezyang.github.io/convolution-visualizer/index.html

GAP(Global average pooling)层

1
gap = torch.nn.AdaptiveAvgPool2d(output_size=1)

双线性汇合(bilinear pooling)

1
2
3
4
5
6
X = torch.reshape(N, D, H * W)                        # Assume X has shape N*D*H*W
X = torch.bmm(X, torch.transpose(X, 1, 2)) / (H * W) # Bilinear pooling
assert X.size() == (N, D, D)
X = torch.reshape(X, (N, D * D))
X = torch.sign(X) * torch.sqrt(torch.abs(X) + 1e-5) # Signed-sqrt normalization
X = torch.nn.functional.normalize(X) # L2 normalization

多卡同步

当使用 torch.nn.DataParallel 将代码运行在多张 GPU 卡上时,PyTorch 的 BN 层默认操作是各卡上数据独立地计算均值和标准差,同步 BN 使用所有卡上的数据一起计算 BN 层的均值和标准差,缓解了当批量大小(batch size)比较小时对均值和标准差估计不准的情况,是在目标检测等任务中一个有效的提升性能的技巧。

链接:https://github.com/vacancy/Synchronized-BatchNorm-PyTorch

类似 BN 滑动平均

如果要实现类似 BN 滑动平均的操作,在 forward 函数中要使用原地(inplace)操作给滑动平均赋值。

1
2
3
4
5
6
7
8
class BN(torch.nn.Module)
def __init__(self):
...
self.register_buffer('running_mean', torch.zeros(num_features))

def forward(self, X):
...
self.running_mean += momentum * (current - self.running_mean)

计算模型整体参数量

1
num_parameters = sum(torch.numel(parameter) for parameter in model.parameters())

类似 Keras 的 model.summary() 输出模型信息

链接:https://github.com/sksq96/pytorch-summary

模型权值初始化

注意 model.modules() 和 model.children() 的区别:model.modules() 会迭代地遍历模型的所有子层,而 model.children() 只会遍历模型下的一层。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Common practise for initialization.
for layer in model.modules():
if isinstance(layer, torch.nn.Conv2d):
torch.nn.init.kaiming_normal_(layer.weight, mode='fan_out',
nonlinearity='relu')
if layer.bias is not None:
torch.nn.init.constant_(layer.bias, val=0.0)
elif isinstance(layer, torch.nn.BatchNorm2d):
torch.nn.init.constant_(layer.weight, val=1.0)
torch.nn.init.constant_(layer.bias, val=0.0)
elif isinstance(layer, torch.nn.Linear):
torch.nn.init.xavier_normal_(layer.weight)
if layer.bias is not None:
torch.nn.init.constant_(layer.bias, val=0.0)

# Initialization with given tensor.
layer.weight = torch.nn.Parameter(tensor)

部分层使用预训练模型

注意如果保存的模型是 torch.nn.DataParallel,则当前的模型也需要是

1
model.load_state_dict(torch.load('model,pth'), strict=False)

将在 GPU 保存的模型加载到 CPU

1
model.load_state_dict(torch.load('model,pth', map_location='cpu'))

自定义损失函数

即使 PyTorch 已经具有了大量标准损失函数,你有时也可能需要创建自己的损失函数。为了做到这一点,你需要创建一个独立的losses.py文件,并且通过扩展nn.Module创建你的自定义损失函数:

1
2
3
4
5
6
7
8
class CustomLoss(torch.nn.Module):

def __init__(self):
super(CustomLoss,self).__init__()

def forward(self,x,y):
loss = torch.mean((x - y)**2)
return loss

特征提取与微调

特征提取

提取 ImageNet 预训练模型某层的卷积特征

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# VGG-16 relu5-3 feature.
model = torchvision.models.vgg16(pretrained=True).features[:-1]
# VGG-16 pool5 feature.
model = torchvision.models.vgg16(pretrained=True).features
# VGG-16 fc7 feature.
model = torchvision.models.vgg16(pretrained=True)
model.classifier = torch.nn.Sequential(*list(model.classifier.children())[:-3])
# ResNet GAP feature.
model = torchvision.models.resnet18(pretrained=True)
model = torch.nn.Sequential(collections.OrderedDict(
list(model.named_children())[:-1]))

with torch.no_grad():
model.eval()
conv_representation = model(image)

提取 ImageNet 预训练模型多层的卷积特征

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
class FeatureExtractor(torch.nn.Module):
"""Helper class to extract several convolution features from the given
pre-trained model.

Attributes:
_model, torch.nn.Module.
_layers_to_extract, list<str> or set<str>

Example:
>>> model = torchvision.models.resnet152(pretrained=True)
>>> model = torch.nn.Sequential(collections.OrderedDict(
list(model.named_children())[:-1]))
>>> conv_representation = FeatureExtractor(
pretrained_model=model,
layers_to_extract={'layer1', 'layer2', 'layer3', 'layer4'})(image)
"""
def __init__(self, pretrained_model, layers_to_extract):
torch.nn.Module.__init__(self)
self._model = pretrained_model
self._model.eval()
self._layers_to_extract = set(layers_to_extract)

def forward(self, x):
with torch.no_grad():
conv_representation = []
for name, layer in self._model.named_children():
x = layer(x)
if name in self._layers_to_extract:
conv_representation.append(x)
return conv_representation

其他预训练模型

链接:https://github.com/Cadene/pretrained-models.pytorch

微调

微调全连接层

1
2
3
4
5
model = torchvision.models.resnet18(pretrained=True)
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Linear(512, 100) # Replace the last fc layer
optimizer = torch.optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9, weight_decay=1e-4)

以较大学习率微调全连接层,较小学习率微调卷积层

1
2
3
4
5
6
model = torchvision.models.resnet18(pretrained=True)
finetuned_parameters = list(map(id, model.fc.parameters()))
conv_parameters = (p for p in model.parameters() if id(p) not in finetuned_parameters)
parameters = [{'params': conv_parameters, 'lr': 1e-3},
{'params': model.fc.parameters()}]
optimizer = torch.optim.SGD(parameters, lr=1e-2, momentum=0.9, weight_decay=1e-4)

模型训练

常用训练和验证数据预处理

其中 ToTensor 操作会将 PIL.Image 或形状为 H×W×D,数值范围为 [0, 255] 的 np.ndarray 转换为形状为 D×H×W,数值范围为 [0.0, 1.0] 的 torch.Tensor。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
train_transform = torchvision.transforms.Compose([
torchvision.transforms.RandomResizedCrop(size=224,
scale=(0.08, 1.0)),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)),
])
val_transform = torchvision.transforms.Compose([
torchvision.transforms.Resize(224),
torchvision.transforms.CenterCrop(224),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225)),
])

训练基本代码框架

对于训练的最佳代码结构,我们需要使用以下两种模式:

  • 使用 prefetch_generator 中的 BackgroundGenerator 来加载下一个批量数据

  • 使用 tqdm 监控训练过程,并展示计算效率,这能帮助我们找到数据加载流程中的瓶颈

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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
# import statements
import torch
import torch.nn as nn
from torch.utils import data
...

# set flags / seeds
torch.backends.cudnn.benchmark = True
np.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed(1)
...

# Start with main code
if __name__ == '__main__':
# argparse for additional flags for experiment
parser = argparse.ArgumentParser(description="Train a network for ...")
...
opt = parser.parse_args()

# add code for datasets (we always use train and validation/ test set)
data_transforms = transforms.Compose([
transforms.Resize((opt.img_size, opt.img_size)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_dataset = datasets.ImageFolder(
root=os.path.join(opt.path_to_data, "train"),
transform=data_transforms)
train_data_loader = data.DataLoader(train_dataset, ...)

test_dataset = datasets.ImageFolder(
root=os.path.join(opt.path_to_data, "test"),
transform=data_transforms)
test_data_loader = data.DataLoader(test_dataset ...)
...

# instantiate network (which has been imported from *networks.py*)
net = MyNetwork(...)
...

# create losses (criterion in pytorch)
criterion_L1 = torch.nn.L1Loss()
...

# if running on GPU and we want to use cuda move model there
use_cuda = torch.cuda.is_available()
if use_cuda:
net = net.cuda()
...

# create optimizers
optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
...

# load checkpoint if needed/ wanted
start_n_iter = 0
start_epoch = 0
if opt.resume:
ckpt = load_checkpoint(opt.path_to_checkpoint) # custom method for loading last checkpoint
net.load_state_dict(ckpt['net'])
start_epoch = ckpt['epoch']
start_n_iter = ckpt['n_iter']
optim.load_state_dict(ckpt['optim'])
print("last checkpoint restored")
...
# if we want to run experiment on multiple GPUs we move the models there
net = torch.nn.DataParallel(net)
...

# typically we use tensorboardX to keep track of experiments
writer = SummaryWriter(...)
# now we start the main loop
n_iter = start_n_iter
for epoch in range(start_epoch, opt.epochs):
# set models to train mode
net.train()
...
# use prefetch_generator and tqdm for iterating through data
pbar = tqdm(enumerate(BackgroundGenerator(train_data_loader, ...)),
total=len(train_data_loader))
start_time = time.time()
# for loop going through dataset
for i, data in pbar:
# data preparation
img, label = data
if use_cuda:
img = img.cuda()
label = label.cuda()
...
# It's very good practice to keep track of preparation time and computation time using tqdm to find any issues in your dataloader
prepare_time = start_time-time.time()
# forward and backward pass
optim.zero_grad()
...
loss.backward()
optim.step()
...

# udpate tensorboardX
writer.add_scalar(..., n_iter)
...
# compute computation time and *compute_efficiency*
process_time = start_time-time.time()-prepare_time
pbar.set_description("Compute efficiency: {:.2f}, epoch: {}/{}:".format(
process_time/(process_time+prepare_time), epoch, opt.epochs))
start_time = time.time()
# maybe do a test pass every x epochs
if epoch % x == x-1:
# bring models to evaluation mode
net.eval()
...
#do some tests
pbar = tqdm(enumerate(BackgroundGenerator(test_data_loader, ...)),
total=len(test_data_loader))
for i, data in pbar:
...

# save checkpoint if needed
...
1
2
3
4
5
6
7
8
for t in epoch(80):
for images, labels in tqdm.tqdm(train_loader, desc='Epoch %3d' % (t + 1)):
images, labels = images.cuda(), labels.cuda()
scores = model(images)
loss = loss_function(scores, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()

标记平滑(label smoothing)

1
2
3
4
5
6
7
8
9
10
11
12
13
for images, labels in train_loader:
images, labels = images.cuda(), labels.cuda()
N = labels.size(0)
# C is the number of classes.
smoothed_labels = torch.full(size=(N, C), fill_value=0.1 / (C - 1)).cuda()
smoothed_labels.scatter_(dim=1, index=torch.unsqueeze(labels, dim=1), value=0.9)

score = model(images)
log_prob = torch.nn.functional.log_softmax(score, dim=1)
loss = -torch.sum(log_prob * smoothed_labels) / N
optimizer.zero_grad()
loss.backward()
optimizer.step()

Mixup

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
beta_distribution = torch.distributions.beta.Beta(alpha, alpha)
for images, labels in train_loader:
images, labels = images.cuda(), labels.cuda()

# Mixup images.
lambda_ = beta_distribution.sample([]).item()
index = torch.randperm(images.size(0)).cuda()
mixed_images = lambda_ * images + (1 - lambda_) * images[index, :]

# Mixup loss.
scores = model(mixed_images)
loss = (lambda_ * loss_function(scores, labels)
+ (1 - lambda_) * loss_function(scores, labels[index]))

optimizer.zero_grad()
loss.backward()
optimizer.step()

L1 正则化

1
2
3
4
5
l1_regularization = torch.nn.L1Loss(reduction='sum')
loss = ... # Standard cross-entropy loss
for param in model.parameters():
loss += torch.sum(torch.abs(param))
loss.backward()

不对偏置项进行 L2 正则化/权值衰减(weight decay)

1
2
3
4
5
bias_list = (param for name, param in model.named_parameters() if name[-4:] == 'bias')
others_list = (param for name, param in model.named_parameters() if name[-4:] != 'bias')
parameters = [{'parameters': bias_list, 'weight_decay': 0},
{'parameters': others_list}]
optimizer = torch.optim.SGD(parameters, lr=1e-2, momentum=0.9, weight_decay=1e-4)

梯度裁剪(gradient clipping)

1
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=20)

计算 Softmax 输出的准确率

1
2
3
4
score = model(images)
prediction = torch.argmax(score, dim=1)
num_correct = torch.sum(prediction == labels).item()
accuruacy = num_correct / labels.size(0)

可视化模型前馈的计算图

链接:https://github.com/szagoruyko/pytorchviz

可视化学习曲线

有 Facebook 自己开发的 Visdom 和 Tensorboard 两个选择。

https://github.com/facebookresearch/visdom

https://github.com/lanpa/tensorboardX

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# Example using Visdom.
vis = visdom.Visdom(env='Learning curve', use_incoming_socket=False)
assert self._visdom.check_connection()
self._visdom.close()
options = collections.namedtuple('Options', ['loss', 'acc', 'lr'])(
loss={'xlabel': 'Epoch', 'ylabel': 'Loss', 'showlegend': True},
acc={'xlabel': 'Epoch', 'ylabel': 'Accuracy', 'showlegend': True},
lr={'xlabel': 'Epoch', 'ylabel': 'Learning rate', 'showlegend': True})

for t in epoch(80):
tran(...)
val(...)
vis.line(X=torch.Tensor([t + 1]), Y=torch.Tensor([train_loss]),
name='train', win='Loss', update='append', opts=options.loss)
vis.line(X=torch.Tensor([t + 1]), Y=torch.Tensor([val_loss]),
name='val', win='Loss', update='append', opts=options.loss)
vis.line(X=torch.Tensor([t + 1]), Y=torch.Tensor([train_acc]),
name='train', win='Accuracy', update='append', opts=options.acc)
vis.line(X=torch.Tensor([t + 1]), Y=torch.Tensor([val_acc]),
name='val', win='Accuracy', update='append', opts=options.acc)
vis.line(X=torch.Tensor([t + 1]), Y=torch.Tensor([lr]),
win='Learning rate', update='append', opts=options.lr)

得到当前学习率

1
2
3
4
5
6
7
# If there is one global learning rate (which is the common case).
lr = next(iter(optimizer.param_groups))['lr']

# If there are multiple learning rates for different layers.
all_lr = []
for param_group in optimizer.param_groups:
all_lr.append(param_group['lr'])

学习率衰减

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# Reduce learning rate when validation accuarcy plateau.
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=5, verbose=True)
for t in range(0, 80):
train(...); val(...)
scheduler.step(val_acc)

# Cosine annealing learning rate.
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=80)
# Reduce learning rate by 10 at given epochs.
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[50, 70], gamma=0.1)
for t in range(0, 80):
scheduler.step()
train(...); val(...)

# Learning rate warmup by 10 epochs.
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda t: t / 10)
for t in range(0, 10):
scheduler.step()
train(...); val(...)

保存与加载断点

注意为了能够恢复训练,我们需要同时保存模型和优化器的状态,以及当前的训练轮数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Save checkpoint.
is_best = current_acc > best_acc
best_acc = max(best_acc, current_acc)
checkpoint = {
'best_acc': best_acc,
'epoch': t + 1,
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
}
model_path = os.path.join('model', 'checkpoint.pth.tar')
torch.save(checkpoint, model_path)
if is_best:
shutil.copy('checkpoint.pth.tar', model_path)

# Load checkpoint.
if resume:
model_path = os.path.join('model', 'checkpoint.pth.tar')
assert os.path.isfile(model_path)
checkpoint = torch.load(model_path)
best_acc = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
print('Load checkpoint at epoch %d.' % start_epoch)

计算准确率、查准率(precision)、召回率(recall)

1
2
3
4
5
6
7
8
9
10
11
12
# data['label'] and data['prediction'] are groundtruth label and prediction 
# for each image, respectively.
accuracy = np.mean(data['label'] == data['prediction']) * 100

# Compute recision and recall for each class.
for c in range(len(num_classes)):
tp = np.dot((data['label'] == c).astype(int),
(data['prediction'] == c).astype(int))
tp_fp = np.sum(data['prediction'] == c)
tp_fn = np.sum(data['label'] == c)
precision = tp / tp_fp * 100
recall = tp / tp_fn * 100

PyTorch 其他注意事项

模型定义

  • 建议有参数的层和汇合(pooling)层使用 torch.nn 模块定义,激活函数直接使用 torch.nn.functional。torch.nn 模块和 torch.nn.functional 的区别在于,torch.nn 模块在计算时底层调用了 torch.nn.functional,但 torch.nn 模块包括该层参数,还可以应对训练和测试两种网络状态。使用 torch.nn.functional 时要注意网络状态,如
1
2
3
def forward(self, x):
...
x = torch.nn.functional.dropout(x, p=0.5, training=self.training)
  • model(x) 前用 model.train() 和 model.eval() 切换网络状态。

  • 不需要计算梯度的代码块用 with torch.no_grad() 包含起来。model.eval() 和 torch.no_grad() 的区别在于,model.eval() 是将网络切换为测试状态,例如 BN 和随机失活(dropout)在训练和测试阶段使用不同的计算方法。torch.no_grad() 是关闭 PyTorch 张量的自动求导机制,以减少存储使用和加速计算,得到的结果无法进行 loss.backward()。

  • torch.nn.CrossEntropyLoss 的输入不需要经过 Softmax。torch.nn.CrossEntropyLoss 等价于 torch.nn.functional.log_softmax + torch.nn.NLLLoss。

  • loss.backward() 前用 optimizer.zero_grad() 清除累积梯度。optimizer.zero_grad() 和 model.zero_grad() 效果一样。

PyTorch 性能与调试

  • torch.utils.data.DataLoader 中尽量设置 pin_memory=True,对特别小的数据集如 MNIST 设置 pin_memory=False 反而更快一些。num_workers 的设置需要在实验中找到最快的取值。

  • 用 del 及时删除不用的中间变量,节约 GPU 存储。

  • 使用 inplace 操作可节约 GPU 存储,如

1
x = torch.nn.functional.relu(x, inplace=True)
  • 减少 CPU 和 GPU 之间的数据传输。例如如果你想知道一个 epoch 中每个 mini-batch 的 loss 和准确率,先将它们累积在 GPU 中等一个 epoch 结束之后一起传输回 CPU 会比每个 mini-batch 都进行一次 GPU 到 CPU 的传输更快。

  • 使用半精度浮点数 half() 会有一定的速度提升,具体效率依赖于 GPU 型号。需要小心数值精度过低带来的稳定性问题。

  • 时常使用 assert tensor.size() == (N, D, H, W) 作为调试手段,确保张量维度和你设想中一致。

  • 除了标记 y 外,尽量少使用一维张量,使用 n*1 的二维张量代替,可以避免一些意想不到的一维张量计算结果。

  • 统计代码各部分耗时

1
2
3
with torch.autograd.profiler.profile(enabled=True, use_cuda=False) as profile:
...
print(profile)

或者在命令行运行

1
python -m torch.utils.bottleneck main.py

PyTorch 中该做和不该做的

  • 在nn.Module的forward方法中避免使用 Numpy 代码

Numpy 是在 CPU 上运行的,它比 torch 的代码运行得要慢一些。由于 torch 的开发思路与 numpy 相似,所以大多数 Numpy 中的函数已经在 PyTorch 中得到了支持。

  • 将DataLoader从主程序的代码中分离

载入数据的工作流程应该独立于你的主训练程序代码。PyTorch 使用background进程更加高效地载入数据,而不会干扰到主训练进程.

  • 不要在每一步中都记录结果

通常而言,我们要训练我们的模型好几千步。因此,为了减小计算开销,每隔 n 步对损失和其它的计算结果进行记录就足够了。尤其是,在训练过程中将中间结果保存成图像,这种开销是非常大的。

  • 使用命令行参数

使用命令行参数设置代码执行时使用的参数(batch 的大小、学习率等)非常方便。一个简单的实验参数跟踪方法,即直接把从parse_args接收到的字典(dict 数据)打印出来:

1
2
3
# saves arguments to config.txt file
opt = parser.parse_args()with open("config.txt", "w") as f:
f.write(opt.__str__())
  • 如果可能的话,请使用Use .detach()从计算图中释放张量

为了实现自动微分,PyTorch 会跟踪所有涉及张量的操作。请使用.detach()来防止记录不必要的操作。

  • 使用.item()打印出标量张量

你可以直接打印变量。然而,我们建议你使用variable.detach()或variable.item()。在早期版本的 PyTorch(< 0.4)中,你必须使用.data访问变量中的张量值。

  • 使用call方法代替nn.Module中的forward方法

参考资料

-------------本文结束感谢您的阅读-------------
;