PyTorch训练代码模板

从参数定义，到网络模型定义，再到训练步骤，验证步骤，测试步骤，总结了一套较为直观的模板。目录如下：

导入包以及设置随机种子
以类的方式定义超参数
定义自己的模型
定义早停类(此步骤可以省略)
定义自己的数据集Dataset,DataLoader
实例化模型，设置loss，优化器等
开始训练以及调整lr
绘图
预测

一、导入包以及设置随机种子

import numpy as np
import torch
import torch.nn as nn
import numpy as np
import pandas as pd
from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

import random
seed = 42
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)

二、以类的方式定义超参数

class argparse():
    pass

args = argparse()
args.epochs, args.learning_rate, args.patience = [30, 0.001, 4]
args.hidden_size, args.input_size= [40, 30]
args.device, = [torch.device("cuda:0" if torch.cuda.is_available() else "cpu"),]

三、定义自己的模型

class Your_model(nn.Module):
    def __init__(self):
        super(Your_model, self).__init__()
        pass

    def forward(self,x):
        pass
        return x

四、定义早停类(此步骤可以省略)

class EarlyStopping():
    def __init__(self,patience=7,verbose=False,delta=0):
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.val_loss_min = np.Inf
        self.delta = delta
    def __call__(self,val_loss,model,path):
        print("val_loss={}".format(val_loss))
        score = -val_loss
        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(val_loss,model,path)
        elif score < self.best_score+self.delta:
            self.counter+=1
            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter>=self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(val_loss,model,path)
            self.counter = 0
    def save_checkpoint(self,val_loss,model,path):
        if self.verbose:
            print(
                f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')
        torch.save(model.state_dict(), path+'/'+'model_checkpoint.pth')
        self.val_loss_min = val_loss

五、定义自己的数据集Dataset,DataLoader

class Dataset_name(Dataset):
    def __init__(self, flag='train'):
        assert flag in ['train', 'test', 'valid']
        self.flag = flag
        self.__load_data__()

    def __getitem__(self, index):
        pass
    def __len__(self):
        pass

    def __load_data__(self, csv_paths: list):
        pass
        print(
            "train_X.shape:{}\ntrain_Y.shape:{}\nvalid_X.shape:{}\nvalid_Y.shape:{}\n"
            .format(self.train_X.shape, self.train_Y.shape, self.valid_X.shape, self.valid_Y.shape))

train_dataset = Dataset_name(flag='train')
train_dataloader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
valid_dataset = Dataset_name(flag='valid')
valid_dataloader = DataLoader(dataset=valid_dataset, batch_size=64, shuffle=True)

六、实例化模型，设置loss，优化器等

model = Your_model().to(args.device)
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(Your_model.parameters(),lr=args.learning_rate)

train_loss = []
valid_loss = []
train_epochs_loss = []
valid_epochs_loss = []

early_stopping = EarlyStopping(patience=args.patience,verbose=True)

七、开始训练以及调整lr

for epoch in range(args.epochs):
    Your_model.train()
    train_epoch_loss = []
    for idx,(data_x,data_y) in enumerate(train_dataloader,0):
        data_x = data_x.to(torch.float32).to(args.device)
        data_y = data_y.to(torch.float32).to(args.device)
        outputs = Your_model(data_x)
        optimizer.zero_grad()
        loss = criterion(data_y,outputs)
        loss.backward()
        optimizer.step()
        train_epoch_loss.append(loss.item())
        train_loss.append(loss.item())
        if idx%(len(train_dataloader)//2)==0:
            print("epoch={}/{},{}/{}of train, loss={}".format(
                epoch, args.epochs, idx, len(train_dataloader),loss.item()))
    train_epochs_loss.append(np.average(train_epoch_loss))

    #=====================valid============================
    Your_model.eval()
    valid_epoch_loss = []
    for idx,(data_x,data_y) in enumerate(valid_dataloader,0):
        data_x = data_x.to(torch.float32).to(args.device)
        data_y = data_y.to(torch.float32).to(args.device)
        outputs = Your_model(data_x)
        loss = criterion(outputs,data_y)
        valid_epoch_loss.append(loss.item())
        valid_loss.append(loss.item())
    valid_epochs_loss.append(np.average(valid_epoch_loss))
    #==================early stopping======================
    early_stopping(valid_epochs_loss[-1],model=Your_model,path=r'c:\\your_model_to_save')
    if early_stopping.early_stop:
        print("Early stopping")
        break
    #====================adjust lr========================
    lr_adjust = {
            2: 5e-5, 4: 1e-5, 6: 5e-6, 8: 1e-6,
            10: 5e-7, 15: 1e-7, 20: 5e-8
        }
    if epoch in lr_adjust.keys():
        lr = lr_adjust[epoch]
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr
        print('Updating learning rate to {}'.format(lr))

八、绘图

plt.figure(figsize=(12,4))
plt.subplot(121)
plt.plot(train_loss[:])
plt.title("train_loss")
plt.subplot(122)
plt.plot(train_epochs_loss[1:],'-o',label="train_loss")
plt.plot(valid_epochs_loss[1:],'-o',label="valid_loss")
plt.title("epochs_loss")
plt.legend()
plt.show()

九、预测

# 此处可定义一个预测集的Dataloader。也可以直接将你的预测数据reshape,添加batch_size=1
Your_model.eval()
predict = Your_model(data)

PyTorch 常用代码段

1. 基本配置

导入包和版本查询

import torch
import torch.nn as nn
import torchvision
print(torch.__version__)
print(torch.version.cuda)
print(torch.backends.cudnn.version())
print(torch.cuda.get_device_name(0))

可复现性

在硬件设备（CPU、GPU）不同时，完全的可复现性无法保证，即使随机种子相同。但是，在同一个设备上，应该保证可复现性。具体做法是，在程序开始的时候固定torch的随机种子，同时也把numpy的随机种子固定。

np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

显卡设置

如果只需要一张显卡

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

如果需要指定多张显卡，比如0，1号显卡。

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

也可以在命令行运行代码时设置显卡：

CUDA_VISIBLE_DEVICES=0,1 python train.py

清除显存

torch.cuda.empty_cache()

也可以使用在命令行重置GPU的指令

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

2. 张量(Tensor)处理

张量的数据类型

PyTorch有10种CPU张量类型和10种GPU张量类型。

Data type	dtype	CPU tensor	GPU tensor
32-bit floating point	`torch.float32` or `torch.float`	`torch.FloatTensor`	`torch.cuda.FloatTensor`
64-bit floating point	`torch.float64` or `torch.double`	`torch.DoubleTensor`	`torch.cuda.DoubleTensor`
16-bit floating point	`torch.float16` or `torch.half`	`torch.HalfTensor`	`torch.cuda.HalfTensor`
16-bit floating point	`torch.bfloat16`	`torch.BFloat16Tensor`	`torch.cuda.BFloat16Tensor`
32-bit complex	`torch.complex32`
64-bit complex	`torch.complex64`
128-bit complex	`torch.complex128` or `torch.cdouble`
8-bit integer (unsigned)	`torch.uint8`	`torch.ByteTensor`	`torch.cuda.ByteTensor`
8-bit integer (signed)	`torch.int8`	`torch.CharTensor`	`torch.cuda.CharTensor`
16-bit integer (signed)	`torch.int16` or `torch.short`	`torch.ShortTensor`	`torch.cuda.ShortTensor`
32-bit integer (signed)	`torch.int32` or `torch.int`	`torch.IntTensor`	`torch.cuda.IntTensor`
64-bit integer (signed)	`torch.int64` or `torch.long`	`torch.LongTensor`	`torch.cuda.LongTensor`
Boolean	`torch.bool`	`torch.BoolTensor`	`torch.cuda.BoolTensor`
quantized 8-bit integer (unsigned)	`torch.quint8`	`torch.ByteTensor`	/
quantized 8-bit integer (signed)	`torch.qint8`	`torch.CharTensor`	/
quantized 32-bit integer (signed)	`torch.qfint32`	`torch.IntTensor`	/
quantized 4-bit integer (unsigned)	`torch.quint4x2`	`torch.ByteTensor`	/

张量基本信息

tensor = torch.randn(3,4,5)
print(tensor.type())  # 数据类型
print(tensor.size())  # 张量的shape，是个元组
print(tensor.dim())   # 维度的数量

命名张量

张量命名是一个非常有用的方法，这样可以方便地使用维度的名字来做索引或其他操作，大大提高了可读性、易用性，防止出错。

# 在PyTorch 1.3之前，需要使用注释
# Tensor[N, C, H, W]
images = torch.randn(32, 3, 56, 56)
images.sum(dim=1)
images.select(dim=1, index=0)

# PyTorch 1.3之后
NCHW = [‘N’, ‘C’, ‘H’, ‘W’]
images = torch.randn(32, 3, 56, 56, names=NCHW)
images.sum('C')
images.select('C', index=0)
# 也可以这么设置
tensor = torch.rand(3,4,1,2,names=('C', 'N', 'H', 'W'))
# 使用align_to可以对维度方便地排序
tensor = tensor.align_to('N', 'C', 'H', 'W')

数据类型转换

# 设置默认类型，pytorch中的FloatTensor远远快于DoubleTensor
torch.set_default_tensor_type(torch.FloatTensor)

# 类型转换
tensor = tensor.cuda()
tensor = tensor.cpu()
tensor = tensor.float()
tensor = tensor.long()

torch.Tensor与np.ndarray转换

除了CharTensor，其他所有CPU上的张量都支持转换为numpy格式然后再转换回来。

ndarray = tensor.cpu().numpy()
tensor = torch.from_numpy(ndarray).float()
tensor = torch.from_numpy(ndarray.copy()).float() # If ndarray has negative stride.

Torch.tensor与PIL.Image转换

# pytorch中的张量默认采用[N, C, H, W]的顺序，并且数据范围在[0,1]，需要进行转置和规范化
# 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
path = r'./figure.jpg'
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的转换

image = PIL.Image.fromarray(ndarray.astype(np.uint8))

ndarray = np.asarray(PIL.Image.open(path))

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

value = torch.rand(1).item()

张量形变

# 在将卷积层输入全连接层的情况下通常需要对张量做形变处理，
# 相比torch.view，torch.reshape可以自动处理输入张量不连续的情况。
tensor = torch.rand(2,3,4)
shape = (6, 4)
tensor = torch.reshape(tensor, shape)

打乱顺序

tensor = tensor[torch.randperm(tensor.size(0))]  # 打乱第一个维度

水平翻转

# pytorch不支持tensor[::-1]这样的负步长操作，水平翻转可以通过张量索引实现
# 假设张量的维度为[N, D, H, W].
tensor = tensor[:,:,:,torch.arange(tensor.size(3) - 1, -1, -1).long()]

复制张量

# 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个10x5的张量，torch.cat的结果是30x5的张量，
而torch.stack的结果是3x10x5的张量。
'''
tensor = torch.cat(list_of_tensors, dim=0)
tensor = torch.stack(list_of_tensors, dim=0)

将整数标签转为one-hot编码

# pytorch的标记默认从0开始
tensor = torch.tensor([0, 2, 1, 3])
N = tensor.size(0)
num_classes = 4
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())

得到非零元素

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

判断两个张量相等

torch.allclose(tensor1, tensor2)  # float tensor
torch.equal(tensor1, tensor2)     # int tensor

张量扩展

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

矩阵乘法

# Matrix multiplcation: (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

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

利用broadcast机制

dist = torch.sqrt(torch.sum((X1[:,None,:] - X2) ** 2, dim=2))

3. 模型定义和操作

一个简单两层卷积网络的示例

# convolutional neural network (2 convolutional layers)
class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(7*7*32, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

model = ConvNet(num_classes).to(device)

卷积层的计算和展示可以用这个网站辅助。

双线性汇合（bilinear pooling）

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

多卡同步 BN（Batch normalization）

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

sync_bn = torch.nn.SyncBatchNorm(num_features, eps=1e-05, momentum=0.1, affine=True, 
                                 track_running_stats=True)

将已有网络的所有BN层改为同步BN层

def convertBNtoSyncBN(module, process_group=None):
    '''Recursively replace all BN layers to SyncBN layer.

    Args:
        module[torch.nn.Module]. Network
    '''
    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
        sync_bn = torch.nn.SyncBatchNorm(module.num_features, module.eps, module.momentum, 
                                         module.affine, module.track_running_stats, process_group)
        sync_bn.running_mean = module.running_mean
        sync_bn.running_var = module.running_var
        if module.affine:
            sync_bn.weight = module.weight.clone().detach()
            sync_bn.bias = module.bias.clone().detach()
        return sync_bn
    else:
        for name, child_module in module.named_children():
            setattr(module, name) = convert_syncbn_model(child_module, process_group=process_group))
        return module

类似 BN 滑动平均

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

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)

计算模型整体参数量

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

查看网络中的参数

可以通过model.state_dict()或者model.named_parameters()函数查看现在的全部可训练参数（包括通过继承得到的父类中的参数）

params = list(model.named_parameters())
(name, param) = params[28]
print(name)
print(param.grad)
print('-------------------------------------------------')
(name2, param2) = params[29]
print(name2)
print(param2.grad)
print('----------------------------------------------------')
(name1, param1) = params[30]
print(name1)
print(param1.grad)

模型可视化（使用pytorchviz）

szagoruyko/pytorchviz

类似 Keras 的 model.summary() 输出模型信息（使用pytorch-summary ）

sksq96/pytorch-summary

模型权重初始化

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

# 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)

提取模型中的某一层

modules()会返回模型中所有模块的迭代器，它能够访问到最内层，比如self.layer1.conv1这个模块，还有一个与它们相对应的是name_children()属性以及named_modules(),这两个不仅会返回模块的迭代器，还会返回网络层的名字。

# 取模型中的前两层
new_model = nn.Sequential(*list(model.children())[:2] 
# 如果希望提取出模型中的所有卷积层，可以像下面这样操作：
for layer in model.named_modules():
    if isinstance(layer[1],nn.Conv2d):
         conv_model.add_module(layer[0],layer[1])

部分层使用预训练模型

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

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

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

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

4. 数据处理

计算数据集的均值和标准差

import os
import cv2
import numpy as np
from torch.utils.data import Dataset
from PIL import Image

def compute_mean_and_std(dataset):
    # 输入PyTorch的dataset，输出均值和标准差
    mean_r = 0
    mean_g = 0
    mean_b = 0

    for img, _ in dataset:
        img = np.asarray(img) # change PIL Image to numpy array
        mean_b += np.mean(img[:, :, 0])
        mean_g += np.mean(img[:, :, 1])
        mean_r += np.mean(img[:, :, 2])

    mean_b /= len(dataset)
    mean_g /= len(dataset)
    mean_r /= len(dataset)

    diff_r = 0
    diff_g = 0
    diff_b = 0

    N = 0

    for img, _ in dataset:
        img = np.asarray(img)

        diff_b += np.sum(np.power(img[:, :, 0] - mean_b, 2))
        diff_g += np.sum(np.power(img[:, :, 1] - mean_g, 2))
        diff_r += np.sum(np.power(img[:, :, 2] - mean_r, 2))

        N += np.prod(img[:, :, 0].shape)

    std_b = np.sqrt(diff_b / N)
    std_g = np.sqrt(diff_g / N)
    std_r = np.sqrt(diff_r / N)

    mean = (mean_b.item() / 255.0, mean_g.item() / 255.0, mean_r.item() / 255.0)
    std = (std_b.item() / 255.0, std_g.item() / 255.0, std_r.item() / 255.0)
    return mean, std

得到视频数据基本信息

import cv2
video = cv2.VideoCapture(mp4_path)
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
fps = int(video.get(cv2.CAP_PROP_FPS))
video.release()

TSN 每段（segment）采样一帧视频

K = self._num_segments
if is_train:
    if num_frames > K:
        # Random index for each segment.
        frame_indices = torch.randint(
            high=num_frames // K, size=(K,), dtype=torch.long)
        frame_indices += num_frames // K * torch.arange(K)
    else:
        frame_indices = torch.randint(
            high=num_frames, size=(K - num_frames,), dtype=torch.long)
        frame_indices = torch.sort(torch.cat((
            torch.arange(num_frames), frame_indices)))[0]
else:
    if num_frames > K:
        # Middle index for each segment.
        frame_indices = num_frames / K // 2
        frame_indices += num_frames // K * torch.arange(K)
    else:
        frame_indices = torch.sort(torch.cat((                              
            torch.arange(num_frames), torch.arange(K - num_frames))))[0]
assert frame_indices.size() == (K,)
return [frame_indices[i] for i in range(K)]

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

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

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(256),
    torchvision.transforms.CenterCrop(224),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
                                     std=(0.229, 0.224, 0.225)),
])

5. 模型训练和测试

分类模型训练代码

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# Train the model
total_step = len(train_loader)
for epoch in range(num_epochs):
    for i ,(images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward and optimizer
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 100 == 0:
            print('Epoch: [{}/{}], Step: [{}/{}], Loss: {}'
                  .format(epoch+1, num_epochs, i+1, total_step, loss.item()))

分类模型测试代码

# Test the model
model.eval()  # eval mode(batch norm uses moving mean/variance 
              #instead of mini-batch mean/variance)
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Test accuracy of the model on the 10000 test images: {} %'
          .format(100 * correct / total))

自定义loss

继承torch.nn.Module类写自己的loss。

pythonclass MyLoss(torch.nn.Moudle):
    def __init__(self):
        super(MyLoss, self).__init__()

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

标签平滑（label smoothing）

写一个label_smoothing.py的文件，然后在训练代码里引用，用LSR代替交叉熵损失即可。label_smoothing.py内容如下：

import torch
import torch.nn as nn

class LSR(nn.Module):

    def __init__(self, e=0.1, reduction='mean'):
        super().__init__()

        self.log_softmax = nn.LogSoftmax(dim=1)
        self.e = e
        self.reduction = reduction

    def _one_hot(self, labels, classes, value=1):
        """
            Convert labels to one hot vectors

        Args:
            labels: torch tensor in format [label1, label2, label3, ...]
            classes: int, number of classes
            value: label value in one hot vector, default to 1

        Returns:
            return one hot format labels in shape [batchsize, classes]
        """

        one_hot = torch.zeros(labels.size(0), classes)

        #labels and value_added  size must match
        labels = labels.view(labels.size(0), -1)
        value_added = torch.Tensor(labels.size(0), 1).fill_(value)

        value_added = value_added.to(labels.device)
        one_hot = one_hot.to(labels.device)

        one_hot.scatter_add_(1, labels, value_added)

        return one_hot

    def _smooth_label(self, target, length, smooth_factor):
        """convert targets to one-hot format, and smooth
        them.
        Args:
            target: target in form with [label1, label2, label_batchsize]
            length: length of one-hot format(number of classes)
            smooth_factor: smooth factor for label smooth

        Returns:
            smoothed labels in one hot format
        """
        one_hot = self._one_hot(target, length, value=1 - smooth_factor)
        one_hot += smooth_factor / (length - 1)

        return one_hot.to(target.device)

    def forward(self, x, target):

        if x.size(0) != target.size(0):
            raise ValueError('Expected input batchsize ({}) to match target batch_size({})'
                    .format(x.size(0), target.size(0)))

        if x.dim() < 2:
            raise ValueError('Expected input tensor to have least 2 dimensions(got {})'
                    .format(x.size(0)))

        if x.dim() != 2:
            raise ValueError('Only 2 dimension tensor are implemented, (got {})'
                    .format(x.size()))

        smoothed_target = self._smooth_label(target, x.size(1), self.e)
        x = self.log_softmax(x)
        loss = torch.sum(- x * smoothed_target, dim=1)

        if self.reduction == 'none':
            return loss

        elif self.reduction == 'sum':
            return torch.sum(loss)

        elif self.reduction == 'mean':
            return torch.mean(loss)

        else:
            raise ValueError('unrecognized option, expect reduction to be one of none, mean, sum')

或者直接在训练文件里做label smoothing

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训练

beta_distribution = torch.distributions.beta.Beta(alpha, alpha)
for images, labels in train_loader:
    images, labels = images.cuda(), labels.cuda()

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

    # Mixup loss.
    scores = model(mixed_images)
    loss = (lambda_ * loss_function(scores, label_a)
            + (1 - lambda_) * loss_function(scores, label_b))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

L1 正则化

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()

不对偏置项进行权重衰减（weight decay）

pytorch里的weight decay相当于l2正则

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）

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

得到当前学习率

# 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'])

另一种方法，在一个batch训练代码里，当前的lr是optimizer.param_groups[0][‘lr’]

学习率衰减

# 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.4版本开始，torch.optim.lr_scheduler 支持链式更新（chaining），即用户可以定义两个 schedulers，并交替在训练中使用。

import torch
from torch.optim import SGD
from torch.optim.lr_scheduler import ExponentialLR, StepLR
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
optimizer = SGD(model, 0.1)
scheduler1 = ExponentialLR(optimizer, gamma=0.9)
scheduler2 = StepLR(optimizer, step_size=3, gamma=0.1)
for epoch in range(4):
    print(epoch, scheduler2.get_last_lr()[0])
    optimizer.step()
    scheduler1.step()
    scheduler2.step()

模型训练可视化

PyTorch可以使用tensorboard来可视化训练过程。

安装和运行TensorBoard。

pip install tensorboard
tensorboard --logdir=runs

使用SummaryWriter类来收集和可视化相应的数据，放了方便查看，可以使用不同的文件夹，比如’Loss/train’和’Loss/test’。

from torch.utils.tensorboard import SummaryWriter
import numpy as np

writer = SummaryWriter()

for n_iter in range(100):
    writer.add_scalar('Loss/train', np.random.random(), n_iter)
    writer.add_scalar('Loss/test', np.random.random(), n_iter)
    writer.add_scalar('Accuracy/train', np.random.random(), n_iter)
    writer.add_scalar('Accuracy/test', np.random.random(), n_iter)

保存与加载断点

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

start_epoch = 0
# Load checkpoint.
if resume: # resume为参数，第一次训练时设为0，中断再训练时设为1
    model_path = os.path.join('model', 'best_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 {}.'.format(start_epoch))
    print('Best accuracy so far {}.'.format(best_acc))

# Train the model
for epoch in range(start_epoch, num_epochs): 
    ... 

    # Test the model
    ...

    # save checkpoint
    is_best = current_acc > best_acc
    best_acc = max(current_acc, best_acc)
    checkpoint = {
        'best_acc': best_acc,
        'epoch': epoch + 1,
        'model': model.state_dict(),
        'optimizer': optimizer.state_dict(),
    }
    model_path = os.path.join('model', 'checkpoint.pth.tar')
    best_model_path = os.path.join('model', 'best_checkpoint.pth.tar')
    torch.save(checkpoint, model_path)
    if is_best:
        shutil.copy(model_path, best_model_path)

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

# 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 预训练模型多层的卷积特征

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

微调全连接层

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)

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

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)

6. 其他注意事项

不要使用太大的线性层。因为nn.Linear(m,n)使用的是的内存，线性层太大很容易超出现有显存。
不要在太长的序列上使用RNN。因为RNN反向传播使用的是BPTT算法，其需要的内存和输入序列的长度呈线性关系。
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()。
model.zero_grad()会把整个模型的参数的梯度都归零, 而optimizer.zero_grad()只会把传入其中的参数的梯度归零.
torch.nn.CrossEntropyLoss 的输入不需要经过 Softmax。torch.nn.CrossEntropyLoss 等价于 torch.nn.functional.log_softmax + torch.nn.NLLLoss。
loss.backward() 前用 optimizer.zero_grad() 清除累积梯度。
torch.utils.data.DataLoader 中尽量设置 pin_memory=True，对特别小的数据集如 MNIST 设置 pin_memory=False 反而更快一些。num_workers 的设置需要在实验中找到最快的取值。
用 del 及时删除不用的中间变量，节约 GPU 存储。
使用 inplace 操作可节约 GPU 存储，如
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 的二维张量代替，可以避免一些意想不到的一维张量计算结果。
统计代码各部分耗时
with torch.autograd.profiler.profile(enabled=True, use_cuda=False) as profile:
…
print(profile)
或者在命令行运行
python -m torch.utils.bottleneck main.py
使用TorchSnooper来调试PyTorch代码，程序在执行的时候，就会自动 print 出来每一行的执行结果的 tensor 的形状、数据类型、设备、是否需要梯度的信息。
pip install torchsnooper
import torchsnooper
对于函数，使用修饰器
@torchsnooper.snoop()
如果不是函数，使用 with 语句来激活 TorchSnooper，把训练的那个循环装进 with 语句中去。
with torchsnooper.snoop():
原本的代码

https://github.com/zasdfgbnm/TorchSnooper

模型可解释性，使用captum库

https://captum.ai/

PyTorch使用技巧

指定GPU编号

设置当前使用的GPU设备仅为0号设备，设备名称为 /gpu:0：os.environ["CUDA_VISIBLE_DEVICES"] = "0"
设置当前使用的GPU设备为0,1号两个设备，名称依次为 /gpu:0、/gpu:1：os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" ，根据顺序表示优先使用0号设备,然后使用1号设备。

指定GPU的命令需要放在和神经网络相关的一系列操作的前面。

查看模型每层输出详情

Keras有一个简洁的API来查看模型的每一层输出尺寸，这在调试网络时非常有用。现在在PyTorch中也可以实现这个功能。

使用很简单，如下用法：

from torchsummary import summary
summary(your_model, input_size=(channels, H, W))

input_size 是根据你自己的网络模型的输入尺寸进行设置。

梯度裁剪（Gradient Clipping）

import torch.nn as nn

outputs = model(data)
loss= loss_fn(outputs, target)
optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), max_norm=20, norm_type=2)
optimizer.step()

nn.utils.clip_grad_norm_ 的参数：

parameters – 一个基于变量的迭代器，会进行梯度归一化
max_norm – 梯度的最大范数
norm_type – 规定范数的类型，默认为L2

扩展单张图片维度

因为在训练时的数据维度一般都是 (batch_size, c, h, w)，而在测试时只输入一张图片，所以需要扩展维度，扩展维度有多个方法：

import cv2
import torch

image = cv2.imread(img_path)
image = torch.tensor(image)
print(image.size())
img = image.view(1, *image.size())
print(img.size())
# output:
# torch.Size([h, w, c])
# torch.Size([1, h, w, c])

或

import cv2
import numpy as np
image = cv2.imread(img_path)
print(image.shape)
img = image[np.newaxis, :, :, :]
print(img.shape)
# output:
# (h, w, c)
# (1, h, w, c)

或

import cv2
import torch

image = cv2.imread(img_path)
image = torch.tensor(image)
print(image.size())

img = image.unsqueeze(dim=0)  
print(img.size())

img = img.squeeze(dim=0)
print(img.size())
# output:
# torch.Size([(h, w, c)])
# torch.Size([1, h, w, c])
# torch.Size([h, w, c])

tensor.unsqueeze(dim)：扩展维度，dim指定扩展哪个维度。

tensor.squeeze(dim)：去除dim指定的且size为1的维度，维度大于1时，squeeze()不起作用，不指定dim时，去除所有size为1的维度。

独热编码

在PyTorch中使用交叉熵损失函数的时候会自动把label转化成onehot，所以不用手动转化，而使用MSE需要手动转化成onehot编码。

import torch
class_num = 8
batch_size = 4
def one_hot(label):    
    """    将一维列表转换为独热编码    """    
    label = label.resize_(batch_size, 1)    
    m_zeros = torch.zeros(batch_size, class_num)    
    # 从 value 中取值，然后根据 dim 和 index 给相应位置赋值    
    onehot = m_zeros.scatter_(1, label, 1)  # (dim,index,value)
    return onehot.numpy()  # Tensor -> Numpy

label = torch.LongTensor(batch_size).random_() % class_num  # 对随机数取余print(one_hot(label))
# output:
[[0. 0. 0. 1. 0. 0. 0. 0.] 
 [0. 0. 0. 0. 1. 0. 0. 0.] 
 [0. 0. 1. 0. 0. 0. 0. 0.] 
 [0. 1. 0. 0. 0. 0. 0. 0.]]

防止验证模型时爆显存

验证模型时不需要求导，即不需要梯度计算，关闭autograd，可以提高速度，节约内存。如果不关闭可能会爆显存。

with torch.no_grad():    
    # 使用model进行预测的代码    
    pass

Pytorch 训练时无用的临时变量可能会越来越多，导致 out of memory ，可以使用下面语句来清理这些不需要的变量。

官网上的解释为：

Releases all unoccupied cached memory currently held by the caching allocator so that those can be used in other GPU application and visible innvidia-smi. torch.cuda.empty_cache()

意思就是PyTorch的缓存分配器会事先分配一些固定的显存，即使实际上tensors并没有使用完这些显存，这些显存也不能被其他应用使用。这个分配过程由第一次CUDA内存访问触发的。

而 torch.cuda.empty_cache() 的作用就是释放缓存分配器当前持有的且未占用的缓存显存，以便这些显存可以被其他GPU应用程序中使用，并且通过 nvidia-smi命令可见。注意使用此命令不会释放tensors占用的显存。

对于不用的数据变量，Pytorch 可以自动进行回收从而释放相应的显存。

更详细的优化可以查看优化显存使用和显存利用问题。

学习率衰减

import torch.optim as optim
from torch.optim import lr_scheduler

# 训练前的初始化
optimizer = optim.Adam(net.parameters(), lr=0.001)
scheduler = lr_scheduler.StepLR(optimizer, 10, 0.1)  # 每过10个epoch，学习率乘以0.1
# 训练过程中
for n in n_epoch:    
    scheduler.step()    
    ...

可以随时查看学习率的值：optimizer.param_groups[0]['lr']。

还有其他学习率更新的方式：

1、自定义更新公式：

scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda epoch:1/(epoch+1))

2、不依赖epoch更新学习率：

lr_scheduler.ReduceLROnPlateau()提供了基于训练中某些测量值使学习率动态下降的方法，它的参数说明到处都可以查到。
提醒一点就是参数 mode=’min’ 还是’max’，取决于优化的的损失还是准确率，即使用 scheduler.step(loss)还是scheduler.step(acc) 。

冻结某些层的参数

参考：https://www.zhihu.com/question/311095447/answer/589307812

在加载预训练模型的时候，我们有时想冻结前面几层，使其参数在训练过程中不发生变化。

我们需要先知道每一层的名字，通过如下代码打印：

net = Network()  # 获取自定义网络结构
for name, value in net.named_parameters():    
    print('name: {0},\t grad: {1}'.format(name, value.requires_grad))

假设前几层信息如下：

name: cnn.VGG_16.convolution1_1.weight, grad: True
name: cnn.VGG_16.convolution1_1.bias, grad: True
name: cnn.VGG_16.convolution1_2.weight, grad: True
name: cnn.VGG_16.convolution1_2.bias, grad: True
name: cnn.VGG_16.convolution2_1.weight, grad: True
name: cnn.VGG_16.convolution2_1.bias, grad: True
name: cnn.VGG_16.convolution2_2.weight, grad: True
name: cnn.VGG_16.convolution2_2.bias, grad: True

后面的True表示该层的参数可训练，然后我们定义一个要冻结的层的列表：

no_grad = [    
    'cnn.VGG_16.convolution1_1.weight',    
    'cnn.VGG_16.convolution1_1.bias',    
    'cnn.VGG_16.convolution1_2.weight',    
    'cnn.VGG_16.convolution1_2.bias']

冻结方法如下：

net = Net.CTPN()  # 获取网络结构
for name, value in net.named_parameters():    
    if name in no_grad:        
        value.requires_grad = False    
    else:        
        value.requires_grad = True

冻结后我们再打印每层的信息：

name: cnn.VGG_16.convolution1_1.weight, grad: False
name: cnn.VGG_16.convolution1_1.bias, grad: False
name: cnn.VGG_16.convolution1_2.weight, grad: False
name: cnn.VGG_16.convolution1_2.bias, grad: False
name: cnn.VGG_16.convolution2_1.weight, grad: True
name: cnn.VGG_16.convolution2_1.bias, grad: True
name: cnn.VGG_16.convolution2_2.weight, grad: True
name: cnn.VGG_16.convolution2_2.bias, grad: True

可以看到前两层的weight和bias的requires_grad都为False，表示它们不可训练。

最后在定义优化器时，只对requires_grad为True的层的参数进行更新。

optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=0.01)

对不同层使用不同学习率

我们对模型的不同层使用不同的学习率。

还是使用这个模型作为例子：

net = Network()  # 获取自定义网络结构
for name, value in net.named_parameters():    
    print('name: {}'.format(name))

# 输出：
# name: cnn.VGG_16.convolution1_1.weight
# name: cnn.VGG_16.convolution1_1.bias
# name: cnn.VGG_16.convolution1_2.weight
# name: cnn.VGG_16.convolution1_2.bias
# name: cnn.VGG_16.convolution2_1.weight
# name: cnn.VGG_16.convolution2_1.bias
# name: cnn.VGG_16.convolution2_2.weight
# name: cnn.VGG_16.convolution2_2.bias

对 convolution1 和 convolution2 设置不同的学习率，首先将它们分开，即放到不同的列表里：

conv1_params = []
conv2_params = []

for name, parms in net.named_parameters():    
    if "convolution1" in name:        
        conv1_params += [parms]    
    else:        
        conv2_params += [parms]

# 然后在优化器中进行如下操作：
optimizer = optim.Adam(    
    [        
        {"params": conv1_params, 'lr': 0.01},        
        {"params": conv2_params, 'lr': 0.001},    
    ],    
    weight_decay=1e-3,
)

我们将模型划分为两部分，存放到一个列表里，每部分就对应上面的一个字典，在字典里设置不同的学习率。当这两部分有相同的其他参数时，就将该参数放到列表外面作为全局参数，如上面的weight_decay。

也可以在列表外设置一个全局学习率，当各部分字典里设置了局部学习率时，就使用该学习率，否则就使用列表外的全局学习率。

模型相关操作

保存加载模型基本用法

1、保存加载整个模型

保存整个网络模型（网络结构+权重参数）。

torch.save(model, 'net.pkl')

直接加载整个网络模型（可能比较耗时）。

model = torch.load('net.pkl')

2、只保存加载模型参数

只保存模型的权重参数（速度快，占内存少）。

torch.save(model.state_dict(), 'net_params.pkl')

因为我们只保存了模型的参数，所以需要先定义一个网络对象，然后再加载模型参数。

# 构建一个网络结构
model = ClassNet()
# 将模型参数加载到新模型中
state_dict = torch.load('net_params.pkl')
model.load_state_dict(state_dict)

保存模型进行推理测试时，只需保存训练好的模型的权重参数，即推荐第二种方法。

主要用法就是上面这些，接下来讲一下PyTorch中保存加载模型内部的一些原理，以及我们可能会遇到的一些特殊的需求。

保存加载自定义模型

上面保存加载的 net.pkl 其实一个字典，通常包含如下内容：

网络结构：输入尺寸、输出尺寸以及隐藏层信息，以便能够在加载时重建模型。
模型的权重参数：包含各网络层训练后的可学习参数，可以在模型实例上调用 state_dict() 方法来获取，比如前面介绍只保存模型权重参数时用到的 model.state_dict()。
优化器参数：有时保存模型的参数需要稍后接着训练，那么就必须保存优化器的状态和所其使用的超参数，也是在优化器实例上调用 state_dict() 方法来获取这些参数。
其他信息：有时我们需要保存一些其他的信息，比如 epoch，batch_size 等超参数。

知道了这些，那么我们就可以自定义需要保存的内容，比如：

# saving a checkpoint assuming the network class named ClassNet
checkpoint = {'model': ClassNet(),
              'model_state_dict': model.state_dict(),
              'optimizer_state_dict': optimizer.state_dict(),
              'epoch': epoch}

torch.save(checkpoint, 'checkpoint.pkl')

上面的 checkpoint 是个字典，里面有4个键值对，分别表示网络模型的不同信息。

然后我们要加载上面保存的自定义的模型：

def load_checkpoint(filepath):
    checkpoint = torch.load(filepath)
    model = checkpoint['model']  # 提取网络结构
    model.load_state_dict(checkpoint['model_state_dict'])  # 加载网络权重参数
    optimizer = TheOptimizerClass()
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])  # 加载优化器参数

    for parameter in model.parameters():
        parameter.requires_grad = False
    model.eval()

    return model

model = load_checkpoint('checkpoint.pkl')

如果加载模型只是为了进行推理测试，则将每一层的 requires_grad 置为 False，即固定这些权重参数；还需要调用 model.eval() 将模型置为测试模式，主要是将 dropout 和 batch normalization 层进行固定，否则模型的预测结果每次都会不同。

如果希望继续训练，则调用 model.train()，以确保网络模型处于训练模式。

state_dict() 也是一个Python字典对象，model.state_dict() 将每一层的可学习参数映射为参数矩阵，其中只包含具有可学习参数的层(卷积层、全连接层等)。

比如下面这个例子：

# Define model
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 8, 5)
        self.bn = nn.BatchNorm2d(8)
        self.conv2 = nn.Conv2d(8, 16, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.bn(x)
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    # Initialize model
    model = TheModelClass()

    # Initialize optimizer
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    print("Model's state_dict:")
    for param_tensor in model.state_dict():
        print(param_tensor, "\t", model.state_dict()[param_tensor].size())

    print("Optimizer's state_dict:")
    for var_name in optimizer.state_dict():
        print(var_name, "\t", optimizer.state_dict()[var_name])

输出为：

Model's state_dict:
conv1.weight            torch.Size([8, 3, 5, 5])
conv1.bias              torch.Size([8])
bn.weight               torch.Size([8])
bn.bias                 torch.Size([8])
bn.running_mean         torch.Size([8])
bn.running_var          torch.Size([8])
bn.num_batches_tracked  torch.Size([])
conv2.weight            torch.Size([16, 8, 5, 5])
conv2.bias              torch.Size([16])
fc1.weight              torch.Size([120, 400])
fc1.bias                torch.Size([120])
fc2.weight              torch.Size([10, 120])
fc2.bias                torch.Size([10])
Optimizer's state_dict:
state            {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [139805696932024, 139805483616008, 139805483616080, 139805483616152, 139805483616440, 139805483616512, 139805483616584, 139805483616656, 139805483616728, 139805483616800]}]

可以看到 model.state_dict() 保存了卷积层，BatchNorm层和最大池化层的信息；而 optimizer.state_dict() 则保存的优化器的状态和相关的超参数。

跨设备保存加载模型

1、在 CPU 上加载在 GPU 上训练并保存的模型（Save on GPU, Load on CPU）：

device = torch.device('cpu')
model = TheModelClass()
# Load all tensors onto the CPU device
model.load_state_dict(torch.load('net_params.pkl', map_location=device))

map_location：a function, torch.device, string or a dict specifying how to remap storage locations

令 torch.load() 函数的 map_location 参数等于 torch.device('cpu') 即可。这里令 map_location 参数等于 'cpu' 也同样可以。

2、在 GPU 上加载在 GPU 上训练并保存的模型（Save on GPU, Load on GPU）：

device = torch.device("cuda")
model = TheModelClass()
model.load_state_dict(torch.load('net_params.pkl'))
model.to(device)

在这里使用 map_location 参数不起作用，要使用 model.to(torch.device("cuda")) 将模型转换为CUDA优化的模型。

还需要对将要输入模型的数据调用 data = data.to(device)，即将数据从CPU转移到GPU。请注意，调用 my_tensor.to(device) 会返回一个 my_tensor 在 GPU 上的副本，它不会覆盖 my_tensor。因此需要手动覆盖张量：my_tensor = my_tensor.to(device)。

3、在 GPU 上加载在 GPU 上训练并保存的模型（Save on CPU, Load on GPU）

device = torch.device("cuda")
model = TheModelClass()
model.load_state_dict(torch.load('net_params.pkl', map_location="cuda:0"))
model.to(device)

当加载包含GPU tensors的模型时，这些tensors 会被默认加载到GPU上，不过是同一个GPU设备。

当有多个GPU设备时，可以通过将 map_location 设定为 *cuda:device_id* 来指定使用哪一个GPU设备，上面例子是指定编号为0的GPU设备。

其实也可以将 torch.device("cuda") 改为 torch.device("cuda:0") 来指定编号为0的GPU设备。

最后调用 model.to(torch.device('cuda')) 来将模型的tensors转换为 CUDA tensors。

下面是PyTorch官方文档上的用法，可以进行参考：

>>> torch.load('tensors.pt')
# Load all tensors onto the CPU
>>> torch.load('tensors.pt', map_location=torch.device('cpu'))
# Load all tensors onto the CPU, using a function
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# Load all tensors onto GPU 1
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# Map tensors from GPU 1 to GPU 0
>>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

CUDA 的用法

在PyTorch中和GPU相关的几个函数：

import torch

# 判断cuda是否可用；
print(torch.cuda.is_available())

# 获取gpu数量；
print(torch.cuda.device_count())

# 获取gpu名字；
print(torch.cuda.get_device_name(0))

# 返回当前gpu设备索引，默认从0开始；
print(torch.cuda.current_device())

# 查看tensor或者model在哪块GPU上
print(torch.tensor([0]).get_device())

我的电脑输出为：

True
1
GeForce RTX 2080 Ti
0

有时我们需要把数据和模型从cpu移到gpu中，有以下两种方法：

use_cuda = torch.cuda.is_available()

# 方法一：
if use_cuda:
    data = data.cuda()
    model.cuda()

# 方法二：
device = torch.device("cuda" if use_cuda else "cpu")
data = data.to(device)
model.to(device)

个人比较习惯第二种方法，可以少一个 if 语句。而且该方法还可以通过设备号指定使用哪个GPU设备，比如使用0号设备：

device = torch.device("cuda:0" if use_cuda else "cpu")

Pytorch内置one_hot函数

感谢@yangyangyang 补充：Pytorch 1.1后，one_hot可以直接用torch.nn.functional.one_hot。

然后我将Pytorch升级到1.2版本，试用了下 one_hot 函数，确实很方便。

具体用法如下：

import torch.nn.functional as F
import torch
tensor =  torch.arange(0, 5) % 3  # tensor([0, 1, 2, 0, 1])
one_hot = F.one_hot(tensor)

# 输出：
# tensor([[1, 0, 0],
#         [0, 1, 0],
#         [0, 0, 1],
#         [1, 0, 0],
#         [0, 1, 0]])

F.one_hot会自己检测不同类别个数，生成对应独热编码。我们也可以自己指定类别数：

tensor =  torch.arange(0, 5) % 3  # tensor([0, 1, 2, 0, 1])
one_hot = F.one_hot(tensor, num_classes=5)

# 输出：
# tensor([[1, 0, 0, 0, 0],
#         [0, 1, 0, 0, 0],
#         [0, 0, 1, 0, 0],
#         [1, 0, 0, 0, 0],
#         [0, 1, 0, 0, 0]])

升级 Pytorch (cpu版本)的命令：conda install pytorch torchvision \-c pytorch

（希望Pytorch升级不会影响项目代码）

网络参数初始化

神经网络的初始化是训练流程的重要基础环节，会对模型的性能、收敛性、收敛速度等产生重要的影响。

以下介绍两种常用的初始化操作。

(1) 使用pytorch内置的torch.nn.init方法。

常用的初始化操作，例如正态分布、均匀分布、xavier初始化、kaiming初始化等都已经实现，可以直接使用。具体详见PyTorch 中 torch.nn.init 中文文档。

init.xavier_uniform(net1[0].weight)

(2) 对于一些更加灵活的初始化方法，可以借助numpy。

对于自定义的初始化方法，有时tensor的功能不如numpy强大灵活，故可以借助numpy实现初始化方法，再转换到tensor上使用。

for layer in net1.modules():    
    if isinstance(layer, nn.Linear): # 判断是否是线性层        
        param_shape = layer.weight.shape        
        layer.weight.data = torch.from_numpy(np.random.normal(0, 0.5, size=param_shape))         
        # 定义为均值为 0，方差为 0.5 的正态分布

加载内置预训练模型

torchvision.models模块的子模块中包含以下模型：

AlexNet
VGG
ResNet
SqueezeNet
DenseNet

导入这些模型的方法为：

import torchvision.models as models
resnet18 = models.resnet18()
alexnet = models.alexnet()
vgg16 = models.vgg16()

有一个很重要的参数为pretrained，默认为False，表示只导入模型的结构，其中的权重是随机初始化的。

如果pretrained 为 True，表示导入的是在ImageNet数据集上预训练的模型。

import torchvision.models as models
resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
vgg16 = models.vgg16(pretrained=True)

Pytorch编写代码基本步骤思想

分为四大步骤：

1、输入处理模块 (X 输入数据，变成网络能够处理的Tensor类型)

2、模型构建模块 (主要负责从输入的数据，得到预测的y^, 这就是我们经常说的前向过程)

3、定义代价函数和优化器模块 (注意，前向过程只会得到模型预测的结果，并不会自动求导和更新，是由这个模块进行处理)

4、构建训练过程 （迭代训练过程，就是上图表情包的训练迭代过程）

这几个模块分别与上图的数字标号1，2，3，4进行一一对应！

知道了上面的宏观思想之后，后面给出每个模块稍微具体一点的解释和具体一个例子，再帮助大家熟悉对应的代码！

1.数据处理

对于数据处理，最为简单的⽅式就是将数据组织成为⼀个。但许多训练需要⽤到mini-batch，直接组织成Tensor不便于我们操作。pytorch为我们提供了Dataset和Dataloader两个类来方便的构建。

torch.utils.data.Dataset

继承Dataset 类需要override 以下⽅法：

torch.utils.data.DataLoader

torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False)

DataLoader Batch。如果选择shuffle = True，每⼀个epoch 后，mini-Batch batch_size 常⻅的使⽤⽅法如下：

2. 模型构建

所有的模型都需要继承torch.nn.Module ，需要实现以下⽅法：

其中forward() ⽅法是前向传播的过程。在实现模型时，我们不需要考虑反向传播。

3. 定义代价函数和优化器

4、构建训练过程

pytorch的训练循环⼤致如下：

5、总结

13个PyTorch特性

1 DatasetFolder

当学习PyTorch时，人们首先要做的事情之一是实现自己的某种Dataset 。这是一个低级错误，没有必要浪费时间写这样的东西。通常，数据集要么是数据列表(或者是numpy数组)，要么磁盘上的文件。所以，把数据在磁盘上组织好，要比写一个自定义的Dataset来加载某种奇怪的格式更好。

分类器最常见的数据格式之一，是有一个带有子文件夹的目录，子文件夹表示类，子文件夹中的文件表示样本，如下所示。

folder/class_0/file1.txt
folder/class_0/file2.txt
folder/class_0/...

folder/class_1/file3.txt
folder/class_1/file4.txt

folder/class_2/file5.txt
folder/class_2/...

有一个内置的方式来加载这类数据集，不管你的数据是图像，文本文件或其他什么，只要使用’DatasetFolder就可以了。令人惊讶的是，这个类是torchvision包的一部分，而不是核心PyTorch。这个类非常全面，你可以从文件夹中过滤文件，使用自定义代码加载它们，并动态转换原始文件。例子:

from torchvision.datasets import DatasetFolder
from pathlib import Path
# I have text files in this folder
ds = DatasetFolder("/Users/marcin/Dev/tmp/my_text_dataset", 
    loader=lambda path: Path(path).read_text(),
    extensions=(".txt",), #only load .txt files
    transform=lambda text: text[:100], # only take first 100 characters
)

# Everything you need is already there
len(ds), ds.classes, ds.class_to_idx
(20, ['novels', 'thrillers'], {'novels': 0, 'thrillers': 1})

如果你在处理图像，还有一个torchvision.datasets.ImageFolder类，它基于DatasetLoader，它被预先配置为加载图像。

2 尽量少用.to(device)，用zeros_like/ones_like之类的代替

我读过很多来自GitHub仓库的PyTorch代码。最让我恼火的是，几乎在每个repo中都有许多*.to(device)行，它们将数据从CPU或GPU转移到其他地方。这样的语句通常会出现在大量的repos或初学者教程中。我强烈建议尽可能少地实现这类操作，并依赖内置的PyTorch功能自动实现这类操作。到处使用.to(device)通常会导致性能下降，还会出现异常：

Expected object of device type cuda but got device type cpu

显然，有些情况下你无法回避它，但大多数情况(如果不是全部)都在这里。其中一种情况是初始化一个全0或全1的张量，这在深度神经网络计算损失的的时候是经常发生的，模型的输出已经在cuda上了，你需要另外的tensor也是在cuda上，这时，你可以使用*_like操作符：

my_output # on any device, if it's cuda then my_zeros will also be on cudamy_zeros = torch.zeros_like(my_output_from_model)

在内部，PyTorch所做的是调用以下操作：

my_zeros = torch.zeros(my_output.size(), dtype=my_output.dtype, layout=my_output.layout, device=my_output.device)

所以所有的设置都是正确的，这样就减少了代码中出现错误的概率。类似的操作包括：

torch.zeros_like()
torch.ones_like()
torch.rand_like()
torch.randn_like()
torch.randint_like()
torch.empty_like()
torch.full_like()

3 Register Buffer ( nn.Module.register_buffer)

这将是我劝人们不要到处使用 .to(device) 的下一步。有时，你的模型或损失函数需要有预先设置的参数，并在调用forward时使用，例如，它可以是一个“权重”参数，它可以缩放损失或一些固定张量，它不会改变，但每次都使用。对于这种情况，请使用nn.Module.register_buffer 方法，它告诉PyTorch将传递给它的值存储在模块中，并将这些值随模块一起移动。如果你初始化你的模块，然后将它移动到GPU，这些值也会自动移动。此外，如果你保存模块的状态，buffers也会被保存！

一旦注册，这些值就可以在forward函数中访问，就像其他模块的属性一样

from torch import nn
import torch

class ModuleWithCustomValues(nn.Module):
    def __init__(self, weights, alpha):
        super().__init__()
        self.register_buffer("weights", torch.tensor(weights))
        self.register_buffer("alpha", torch.tensor(alpha))

    def forward(self, x):
        return x * self.weights + self.alpha

m = ModuleWithCustomValues(
    weights=[1.0, 2.0], alpha=1e-4
)
m(torch.tensor([1.23, 4.56]))
tensor([1.2301, 9.1201])

4 Built-in Identity()

有时候，当你使用迁移学习时，你需要用1:1的映射替换一些层，可以用nn.Module来实现这个目的，只返回输入值。PyTorch内置了这个类。

例子，你想要在分类层之前从一个预训练过的ResNet50获取图像表示。以下是如何做到这一点：

from torchvision.models import resnet50
model = resnet50(pretrained=True)
model.fc = nn.Identity()
last_layer_output = model(torch.rand((1, 3, 224, 224)))
last_layer_output.shape
torch.Size([1, 2048])

5 Pairwise distances: torch.cdist

下次当你遇到计算两个张量之间的欧几里得距离(或者一般来说:p范数)的问题时，请记住torch.cdist。它确实做到了这一点，并且在使用欧几里得距离时还自动使用矩阵乘法，从而提高了性能。

points1 = torch.tensor([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]])
points2 = torch.tensor([[0.0, 0.0], [-1.0, -1.0], [-2.0, -2.0], [-3.0, -3.0]]) # batches don't have to be equal
torch.cdist(points1, points2, p=2.0)
tensor([[0.0000, 1.4142, 2.8284, 4.2426],
        [1.4142, 2.8284, 4.2426, 5.6569],
        [2.8284, 4.2426, 5.6569, 7.0711]])

没有矩阵乘法或有矩阵乘法的性能，在我的机器上使用mm时，速度快了2倍以上。

%%timeit
points1 = torch.rand((512, 2))
points2 = torch.rand((512, 2))
torch.cdist(points1, points2, p=2.0, compute_mode="donot_use_mm_for_euclid_dist")

867µs±142µs per loop (mean±std. dev. of 7 run, 1000 loop each)

%%timeit
points1 = torch.rand((512, 2))
points2 = torch.rand((512, 2))
torch.cdist(points1, points2, p=2.0)

417µs±52.9µs per loop (mean±std. dev. of 7 run, 1000 loop each)

6 Cosine similarity: F.cosine_similarity

与上一点相同，计算欧几里得距离并不总是你需要的东西。当处理向量时，通常余弦相似度是选择的度量。PyTorch也有一个内置的余弦相似度实现。

import torch.nn.functional as F
vector1 = torch.tensor([0.0, 1.0])
vector2 = torch.tensor([0.05, 1.0])
print(F.cosine_similarity(vector1, vector2, dim=0))
vector3 = torch.tensor([0.0, -1.0])
print(F.cosine_similarity(vector1, vector3, dim=0))
tensor(0.9988)
tensor(-1.)

PyTorch中批量计算余弦距离

import torch.nn.functional as F
batch_of_vectors = torch.rand((4, 64))
similarity_matrix = F.cosine_similarity(batch_of_vectors.unsqueeze(1), batch_of_vectors.unsqueeze(0), dim=2)
similarity_matrix
tensor([[1.0000, 0.6922, 0.6480, 0.6789],
        [0.6922, 1.0000, 0.7143, 0.7172],
        [0.6480, 0.7143, 1.0000, 0.7312],
        [0.6789, 0.7172, 0.7312, 1.0000]])

7 归一化向量: F.normalize

最后一点仍然与向量和距离有松散的联系，那就是归一化：通常是通过改变向量的大小来提高计算的稳定性。最常用的归一化是L2，可以在PyTorch中按如下方式应用:

vector = torch.tensor([99.0, -512.0, 123.0, 0.1, 6.66])
normalized_vector = F.normalize(vector, p=2.0, dim=0)
normalized_vector
tensor([ 1.8476e-01, -9.5552e-01,  2.2955e-01,  1.8662e-04,  1.2429e-02])

在PyTorch中执行归一化的旧方法是：

vector = torch.tensor([99.0, -512.0, 123.0, 0.1, 6.66])
normalized_vector = vector / torch.norm(vector, p=2.0)
normalized_vector
tensor([ 1.8476e-01, -9.5552e-01,  2.2955e-01,  1.8662e-04,  1.2429e-02])

在PyTorch中批量进行L2归一化

batch_of_vectors = torch.rand((4, 64))
normalized_batch_of_vectors = F.normalize(batch_of_vectors, p=2.0, dim=1)
normalized_batch_of_vectors.shape, torch.norm(normalized_batch_of_vectors, dim=1) 
# all vectors will have length of 1.0(torch.Size([4, 64]), tensor([1.0000, 1.0000, 1.0000, 1.0000]))

8 线性层 + 分块技巧 (torch.chunk)

这是我最近发现的一个有创意的技巧。假设你想把你的输入映射到N个不同的线性投影中。你可以通过创建N个nn.Linear来做到这一点。或者你也可以创建一个单一的线性层，做一个向前传递，然后将输出分成N块。这种方法通常会带来更高的性能，所以这是一个值得记住的技巧。

d = 1024
batch = torch.rand((8, d))
layers = nn.Linear(d, 128, bias=False), nn.Linear(d, 128, bias=False), nn.Linear(d, 128, bias=False)
one_layer = nn.Linear(d, 128 * 3, bias=False)
%%timeit
o1 = layers[0](batch)
o2 = layers[1](batch)
o3 = layers[2](batch)

289 µs ± 30.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%%timeit
o1, o2, o3 = torch.chunk(one_layer(batch), 3, dim=1)

202 µs ± 8.09 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

9 Masked select (torch.masked_select)

有时你只需要对输入张量的一部分进行计算。给你一个例子：你想计算的损失只在满足某些条件的张量上。为了做到这一点，你可以使用torch.masked_select，注意，当需要梯度时也可以使用这个操作。

data = torch.rand((3, 3)).requires_grad_()
print(data)
mask = data > data.mean()
print(mask)
torch.masked_select(data, mask)
tensor([[0.0582, 0.7170, 0.7713],
        [0.9458, 0.2597, 0.6711],
        [0.2828, 0.2232, 0.1981]], requires_grad=True)
tensor([[False,  True,  True],
        [ True, False,  True],
        [False, False, False]])
tensor([0.7170, 0.7713, 0.9458, 0.6711], grad_fn=<MaskedSelectBackward>)

直接在tensor上应用mask

类似的行为可以通过使用mask作为输入张量的 “indexer”来实现。

data[mask]
tensor([0.7170, 0.7713, 0.9458, 0.6711], grad_fn=<IndexBackward>)

有时，一个理想的解决方案是用0填充mask中所有的False值，可以这样做:

data * mask
tensor([[0.0000, 0.7170, 0.7713],
        [0.9458, 0.0000, 0.6711],
        [0.0000, 0.0000, 0.0000]], grad_fn=<MulBackward0>)

10 使用 torch.where来对tensors加条件

当你想把两个张量结合在一个条件下这个函数很有用，如果条件是真，那么从第一个张量中取元素，如果条件是假，从第二个张量中取元素。

x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0], requires_grad=True)
y = -x
condition_or_mask = x <= 3.0
torch.where(condition_or_mask, x, y)
tensor([ 1.,  2.,  3., -4., -5.], grad_fn=<SWhereBackward>)

11 在给定的位置给张量填入值(Tensor.scatter)

这个函数的用例如下，你想用给定位置下另一个张量的值填充一个张量。一维张量更容易理解，所以我将先展示它，然后继续更高级的例子。

data = torch.tensor([1, 2, 3, 4, 5])
index = torch.tensor([0, 1])
values = torch.tensor([-1, -2, -3, -4, -5])
data.scatter(0, index, values)
tensor([-1, -2,  3,  4,  5])

上面的例子很简单，但是现在看看如果将index改为index = torch.tensor([0, 1, 4])会发生什么：

data = torch.tensor([1, 2, 3, 4, 5])
index = torch.tensor([0, 1, 4])
values = torch.tensor([-1, -2, -3, -4, -5])
data.scatter(0, index, values)
tensor([-1, -2,  3,  4, -3])

为什么最后一个值是-3，这是反直觉的，对吧？这是PyTorch scatter函数的中心思想。index变量表示data张量的第i个值应该放在values张量的哪个位置。我希望下面的简单python版的这个操作能让你更明白：

data_orig = torch.tensor([1, 2, 3, 4, 5])
index = torch.tensor([0, 1, 4])
values = torch.tensor([-1, -2, -3, -4, -5])
scattered = data_orig.scatter(0, index, values)

data = data_orig.clone()
for idx_in_values, where_to_put_the_value in enumerate(index):
    what_value_to_put = values[idx_in_values]
    data[where_to_put_the_value] = what_value_to_put
data, scattered
(tensor([-1, -2,  3,  4, -3]), tensor([-1, -2,  3,  4, -3]))

2D数据的PyTorch scatter例子

始终记住，index的形状与values的形状相关，而index中的值对应于data中的位置。

data = torch.zeros((4, 4)).float()
index = torch.tensor([
    [0, 1],
    [2, 3],
    [0, 3],
    [1, 2]
])
values = torch.arange(1, 9).float().view(4, 2)
values, data.scatter(1, index, values)
(tensor([[1., 2.],
        [3., 4.],
        [5., 6.],
        [7., 8.]]),
tensor([[1., 2., 0., 0.],
        [0., 0., 3., 4.],
        [5., 0., 0., 6.],
        [0., 7., 8., 0.]]))

12 在网络中进行图像插值 (F.interpolate)

当我学习PyTorch时，我惊讶地发现，实际上可以在前向传递中调整图像(或任何中间张量)，并保持梯度流。这种方法在使用CNN和GANs时特别有用。

# image from https://commons.wikimedia.org/wiki/File:A_female_British_Shorthair_at_the_age_of_20_months.jpg
img = Image.open("./cat.jpg")
img

to_pil_image(
    F.interpolate(to_tensor(img).unsqueeze(0),  # batch of size 1
                  mode="bilinear", 
                  scale_factor=2.0, 
                  align_corners=False).squeeze(0) # remove batch dimension
)

看看梯度流是如何保存的：

F.interpolate(to_tensor(img).unsqueeze(0).requires_grad_(),
                  mode="bicubic", 
                  scale_factor=2.0, 
                  align_corners=False)
tensor([[[[0.9216, 0.9216, 0.9216,  ..., 0.8361, 0.8272, 0.8219],
    [0.9214, 0.9214, 0.9214,  ..., 0.8361, 0.8272, 0.8219],
    [0.9212, 0.9212, 0.9212,  ..., 0.8361, 0.8272, 0.8219],
    ...,
    [0.9098, 0.9098, 0.9098,  ..., 0.3592, 0.3486, 0.3421],
    [0.9098, 0.9098, 0.9098,  ..., 0.3566, 0.3463, 0.3400],
    [0.9098, 0.9098, 0.9098,  ..., 0.3550, 0.3449, 0.3387]],

    [[0.6627, 0.6627, 0.6627,  ..., 0.5380, 0.5292, 0.5238],
    [0.6626, 0.6626, 0.6626,  ..., 0.5380, 0.5292, 0.5238],
    [0.6623, 0.6623, 0.6623,  ..., 0.5380, 0.5292, 0.5238],
    ...,
    [0.6196, 0.6196, 0.6196,  ..., 0.3631, 0.3525, 0.3461],
    [0.6196, 0.6196, 0.6196,  ..., 0.3605, 0.3502, 0.3439],
    [0.6196, 0.6196, 0.6196,  ..., 0.3589, 0.3488, 0.3426]],

    [[0.4353, 0.4353, 0.4353,  ..., 0.1913, 0.1835, 0.1787],
    [0.4352, 0.4352, 0.4352,  ..., 0.1913, 0.1835, 0.1787],
    [0.4349, 0.4349, 0.4349,  ..., 0.1913, 0.1835, 0.1787],
    ...,
    [0.3333, 0.3333, 0.3333,  ..., 0.3827, 0.3721, 0.3657],
    [0.3333, 0.3333, 0.3333,  ..., 0.3801, 0.3698, 0.3635],
    [0.3333, 0.3333, 0.3333,  ..., 0.3785, 0.3684, 0.3622]]]],
grad_fn=<UpsampleBicubic2DBackward1>)

13 将图像做成网格 (torchvision.utils.make_grid)

当使用PyTorch和torchvision时，不需要使用matplotlib或一些外部库来复制粘贴代码来显示图像网格。只要使用torchvision.utils.make_grid就行了。

from torchvision.utils import make_grid
from torchvision.transforms.functional import to_tensor, to_pil_image
from PIL import Image
img = Image.open("./cat.jpg")
to_pil_image(
    make_grid(
        [to_tensor(i) for i in [img, img, img]],
         nrow=2, # number of images in single row
         padding=5 # "frame" size
     )
)

深度学习调参技巧

不管什么模型，先在一个较小的训练集上train和test，看看它能不能过拟合。如果不能过拟合，可能是学习率太大，或者代码写错了。先调小学习率试一下，如果还不行就去检查代码，先看dataloader输出的数据对不对，再看模型每一步的size是否符合自己期待。
看train/eval的loss曲线，正常的情况应该是train loss呈log状一直下降最后趋于稳定，eval loss开始时一直下降到某一个epoch之后开始趋于稳定或开始上升，这时候可以用early stopping保存eval loss最低的那个模型。如果loss曲线非常不正常，很有可能是数据处理出了问题，比如label对应错了，回去检查代码。
不要一开始就用大数据集，先在一个大概2w训练集，2k测试集的小数据集上调参。
尽量不要自己从头搭架子（新手和半新手）。找一个已经明确没有bug能跑通的其它任务的架子，在它的基础上修改。否则debug过程非常艰难，因为有时候是版本迭代产生的问题，修改起来很麻烦。
优化器优先用adam，学习率设1e-3或1e-4，再试Radam（LiyuanLucasLiu/RAdam）。不推荐sgdm，因为很慢。
lrscheduler用torch.optim.lr_scheduler.CosineAnnealingLR，T_max设32或64，几个任务上试效果都不错。（用这个lr_scheduler加上adam系的optimizer基本就不用怎么调学习率了）
有一些任务（尤其是有RNN的）要做梯度裁剪，torch.nn.utils.clip_grad_norm。
参数初始化，lstm的h用orthogonal，其它用he或xavier。
激活函数用relu一般就够了，也可以试试leaky relu。
batchnorm和dropout可以试，放的位置很重要。优先尝试放在最后输出层之前，以及embedding层之后。RNN可以试layer_norm。有些任务上加了这些层可能会有负作用。
metric learning中先试标label的分类方法。然后可以用triplet loss，margin这个参数的设置很重要。
batchsize设置小一点通常会有一些提升，某些任务batchsize设成1有奇效。
embedding层的embedsize可以小一些（64 or 128），之后LSTM或CNN的hiddensize要稍微大一些（256 or 512）。（ALBERT论文里面大概也是这个意思）
模型方面，可以先用2或3层LSTM试一下，通常效果都不错。
weight decay可以试一下，我一般用1e-4。
有CNN的地方就用shortcut。CNN层数加到某一个值之后对结果影响就不大了，这个值作为参数可以调一下。
GRU和LSTM在大部分任务上效果差不多。
看论文时候不要全信，能复现的尽量复现一下，许多论文都会做低baseline，但实际使用时很多baseline效果很不错。
对于大多数任务，数据比模型重要。面对新任务时先分析数据，再根据数据设计模型，并决定各个参数。例如nlp有些任务中的padding长度，通常需要达到数据集的90%以上，可用pandas的describe函数进行分析。

PyTorch模型训练提速的技巧

Pytorch-Lightning

你可以在Pytorch的库Pytorch- lightning中找到我在这里讨论的每一个优化。Lightning是在Pytorch之上的一个封装，它可以自动训练，同时让研究人员完全控制关键的模型组件。Lightning 使用最新的最佳实践，并将你可能出错的地方最小化。

我们为MNIST定义LightningModel并使用Trainer来训练模型。

from pytorch_lightning import Trainer
model = LightningModule(…)
trainer = Trainer()
trainer.fit(model)

1. DataLoaders

这可能是最容易获得速度增益的地方。保存h5py或numpy文件以加速数据加载的时代已经一去不复返了，使用Pytorch dataloader加载图像数据很简单(对于NLP数据，请查看TorchText)。

在lightning中，你不需要指定训练循环，只需要定义dataLoaders和Trainer就会在需要的时候调用它们。

dataset = MNIST(root=self.hparams.data_root, train=train, download=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
  x, y = batch
  model.training_step(x, y)
  ...

2. DataLoaders 中的 workers 的数量

另一个加速的神奇之处是允许批量并行加载。因此，您可以一次装载nb_workers个batch，而不是一次装载一个batch。

# slow
loader = DataLoader(dataset, batch_size=32, shuffle=True)
# fast (use 10 workers)
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=10)

3. Batch size

在开始下一个优化步骤之前，将batch size增大到CPU-RAM或GPU-RAM所允许的最大范围。

下一节将重点介绍如何帮助减少内存占用，以便你可以继续增加batch size。

记住，你可能需要再次更新你的学习率。一个好的经验法则是，如果batch size加倍，那么学习率就加倍。

4. 梯度累加

在你已经达到计算资源上限的情况下，你的batch size仍然太小(比如8)，然后我们需要模拟一个更大的batch size来进行梯度下降，以提供一个良好的估计。

假设我们想要达到128的batch size大小。我们需要以batch size为8执行16个前向传播和向后传播，然后再执行一次优化步骤。

# clear last step
optimizer.zero_grad()

# 16 accumulated gradient steps
scaled_loss = 0
for accumulated_step_i in range(16):
     out = model.forward()
     loss = some_loss(out,y)    
     loss.backward()
      scaled_loss += loss.item()

# update weights after 8 steps. effective batch = 8*16
optimizer.step()

# loss is now scaled up by the number of accumulated batches
actual_loss = scaled_loss / 16

在lightning中，全部都给你做好了，只需要设置accumulate_grad_batches=16：

trainer = Trainer(accumulate_grad_batches=16)
trainer.fit(model)

5. 保留的计算图

一个最简单撑爆你的内存的方法是为了记录日志存储你的loss。

losses = []
...
losses.append(loss)

print(f'current loss: {torch.mean(losses)'})

上面的问题是，loss仍然包含有整个图的副本。在这种情况下，调用.item()来释放它。

![1_CER3v8cok2UOBNsmnBrzPQ](9 Tips For Training Lightning-Fast Neural Networks In Pytorch.assets/1_CER3v8cok2UOBNsmnBrzPQ.gif)# bad
losses.append(loss)

# good
losses.append(loss.item())

Lightning会非常小心，确保不会保留计算图的副本。

6. 单个GPU训练

一旦你已经完成了前面的步骤，是时候进入GPU训练了。在GPU上的训练将使多个GPU cores之间的数学计算并行化。你得到的加速取决于你所使用的GPU类型。我推荐个人用2080Ti，公司用V100。

乍一看，这可能会让你不知所措，但你真的只需要做两件事:1)移动你的模型到GPU, 2)每当你运行数据通过它，把数据放到GPU上。

# put model on GPU
model.cuda(0)

# put data on gpu (cuda on a variable returns a cuda copy)
x = x.cuda(0)

# runs on GPU now
model(x)

如果你使用Lightning，你什么都不用做，只需要设置Trainer(gpus=1)。

# ask lightning to use gpu 0 for training
trainer = Trainer(gpus=[0])
trainer.fit(model)

在GPU上进行训练时，要注意的主要事情是限制CPU和GPU之间的传输次数。

# expensive
x = x.cuda(0)# very expensive
x = x.cpu()
x = x.cuda(0)

如果内存耗尽，不要将数据移回CPU以节省内存。在求助于GPU之前，尝试以其他方式优化你的代码或GPU之间的内存分布。

另一件需要注意的事情是调用强制GPU同步的操作。清除内存缓存就是一个例子。

# really bad idea. Stops all the GPUs until they all catch up
torch.cuda.empty_cache()

但是，如果使用Lightning，惟一可能出现问题的地方是在定义Lightning Module时。Lightning会特别注意不去犯这类错误。

7. 16-bit 精度

16bit精度是将内存占用减半的惊人技术。大多数模型使用32bit精度数字进行训练。然而，最近的研究发现，16bit模型也可以工作得很好。混合精度意味着对某些内容使用16bit，但将权重等内容保持在32bit。

要在Pytorch中使用16bit精度，请安装NVIDIA的apex库，并对你的模型进行这些更改。

# enable 16-bit on the model and the optimizer
model, optimizers = amp.initialize(model, optimizers, opt_level='O2')

# when doing .backward, let amp do it so it can scale the loss
with amp.scale_loss(loss, optimizer) as scaled_loss:                      
    scaled_loss.backward()

amp包会处理好大部分事情。如果梯度爆炸或趋向于0，它甚至会缩放loss。

在lightning中，启用16bit并不需要修改模型中的任何内容，也不需要执行我上面所写的操作。设置Trainer(precision=16)就可以了。

trainer = Trainer(amp_level='O2', use_amp=False)
trainer.fit(model)

8. 移动到多个GPUs中

现在，事情变得非常有趣了。有3种(也许更多？)方法来进行多GPU训练。

分batch训练

第一种方法被称为“分batch训练”。该策略将模型复制到每个GPU上，每个GPU获得batch的一部分。

# copy model on each GPU and give a fourth of the batch to each
model = DataParallel(model, devices=[0, 1, 2 ,3])

# out has 4 outputs (one for each gpu)
out = model(x.cuda(0))

在lightning中，你只需要增加GPUs的数量，然后告诉trainer，其他什么都不用做。

# ask lightning to use 4 GPUs for training
trainer = Trainer(gpus=[0, 1, 2, 3])
trainer.fit(model)

模型分布训练

有时你的模型可能太大不能完全放到内存中。例如，带有编码器和解码器的序列到序列模型在生成输出时可能会占用20GB RAM。在本例中，我们希望将编码器和解码器放在独立的GPU上。

# each model is sooo big we can't fit both in memory
encoder_rnn.cuda(0)
decoder_rnn.cuda(1)

# run input through encoder on GPU 0
encoder_out = encoder_rnn(x.cuda(0))

# run output through decoder on the next GPU
out = decoder_rnn(encoder_out.cuda(1))

# normally we want to bring all outputs back to GPU 0
out = out.cuda(0)

对于这种类型的训练，在Lightning中不需要指定任何GPU，你应该把LightningModule中的模块放到正确的GPU上。

class MyModule(LightningModule):
    def __init__():
        self.encoder = RNN(...)
        self.decoder = RNN(...)
    def forward(x):
        # models won't be moved after the first forward because 
        # they are already on the correct GPUs
        self.encoder.cuda(0)
        self.decoder.cuda(1)
        out = self.encoder(x)
        out = self.decoder(out.cuda(1))

# don't pass GPUs to trainer
model = MyModule()
trainer = Trainer()
trainer.fit(model)

两者混合

在上面的情况下，编码器和解码器仍然可以从并行化操作中获益。

# change these lines
self.encoder = RNN(...)
self.decoder = RNN(...)

# to these
# now each RNN is based on a different gpu set
self.encoder = DataParallel(self.encoder, devices=[0, 1, 2, 3])
self.decoder = DataParallel(self.encoder, devices=[4, 5, 6, 7])

# in forward...
out = self.encoder(x.cuda(0))

# notice inputs on first gpu in device
sout = self.decoder(out.cuda(4))  # <--- the 4 here

使用多个GPU时要考虑的注意事项：

如果模型已经在GPU上了，model.cuda()不会做任何事情。
总是把输入放在设备列表中的第一个设备上。
在设备之间传输数据是昂贵的，把它作为最后的手段。
优化器和梯度会被保存在GPU 0上，因此，GPU 0上使用的内存可能会比其他GPU大得多。

9. 多节点GPU训练

每台机器上的每个GPU都有一个模型的副本。每台机器获得数据的一部分，并且只在那部分上训练。每台机器都能同步梯度。

如果你已经做到了这一步，那么你现在可以在几分钟内训练Imagenet了！这并没有你想象的那么难，但是它可能需要你对计算集群的更多知识。这些说明假设你正在集群上使用SLURM。

Pytorch允许多节点训练，通过在每个节点上复制每个GPU上的模型并同步梯度。所以，每个模型都是在每个GPU上独立初始化的，本质上独立地在数据的一个分区上训练，除了它们都从所有模型接收梯度更新。

在高层次上：

在每个GPU上初始化一个模型的副本(确保设置种子，让每个模型初始化到相同的权重，否则它会失败)。
将数据集分割成子集(使用DistributedSampler)。每个GPU只在它自己的小子集上训练。
在.backward()上，所有副本都接收到所有模型的梯度副本。这是模型之间唯一一次的通信。

Pytorch有一个很好的抽象，叫做DistributedDataParallel，它可以帮你实现这个功能。要使用DDP，你需要做4的事情：

def tng_dataloader():
     d = MNIST()

     # 4: Add distributed sampler
     # sampler sends a portion of tng data to each machine
     dist_sampler = DistributedSampler(dataset)
     dataloader = DataLoader(d, shuffle=False, sampler=dist_sampler)

def main_process_entrypoint(gpu_nb):
     # 2: set up connections  between all gpus across all machines
     # all gpus connect to a single GPU "root"
     # the default uses env://
     world = nb_gpus * nb_nodes
     dist.init_process_group("nccl", rank=gpu_nb, world_size=world)

     # 3: wrap model in DPP
     torch.cuda.set_device(gpu_nb)
     model.cuda(gpu_nb)
     model = DistributedDataParallel(model, device_ids=[gpu_nb])

     # train your model now...

if  __name__ == '__main__':
     # 1: spawn number of processes
     # your cluster will call main for each machine
     mp.spawn(main_process_entrypoint, nprocs=8)

然而，在Lightning中，只需设置节点数量，它就会为你处理其余的事情。

# train on 1024 gpus across 128 nodes
trainer = Trainer(nb_gpu_nodes=128, gpus=[0, 1, 2, 3, 4, 5, 6, 7])

Lightning还附带了一个SlurmCluster管理器，可以方便地帮助你提交SLURM作业的正确详细信息。

10. 在单个节点上多GPU更快的训练

事实证明，distributedDataParallel比DataParallel快得多，因为它只执行梯度同步的通信。所以，一个好的hack是使用distributedDataParallel替换DataParallel，即使是在单机上进行训练。

在Lightning中，这很容易通过将distributed_backend设置为ddp和设置GPUs的数量来实现。

# train on 4 gpus on the same machine MUCH faster than DataParallel
trainer = Trainer(distributed_backend='ddp', gpus=[0, 1, 2, 3])

对模型加速的思考

首先，我要确保在数据加载中没有瓶颈。为此，我使用了我所描述的现有数据加载解决方案，但是如果没有一种解决方案满足你的需要，请考虑离线处理和缓存到高性能数据存储中，比如h5py。

接下来看看你在训练步骤中要做什么。确保你的前向传播速度快，避免过多的计算以及最小化CPU和GPU之间的数据传输。最后，避免做一些会降低GPU速度的事情(本指南中有介绍)。

接下来，我试图最大化我的batch size，这通常是受GPU内存大小的限制。现在，需要关注在使用大的batch size的时候如何在多个GPUs上分布并最小化延迟（比如，我可能会尝试着在多个gpu上使用8000 +的有效batch size）。

然而，你需要小心大的batch size。针对你的具体问题，请查阅相关文献，看看人们都忽略了什么！

杰克成

https://jackhcc.github.io/posts/dl-series1.html

本博客所有文章除特別声明外，均采用 CC BY 4.0 许可协议。转载请注明来源杰克成 !

PyTorch

DL专栏2-Activate Function

常见的激活函数

2021-10-05 Deep Learning

Activate Function

Python Module

Python好用的小众模块

2021-10-03 Python

Python

DL专栏1-PyTorch Tricks

PyTorch训练代码模板

一、导入包以及设置随机种子

二、以类的方式定义超参数

三、定义自己的模型

四、定义早停类(此步骤可以省略)

五、定义自己的数据集Dataset,DataLoader

六、实例化模型，设置loss，优化器等

七、开始训练以及调整lr

八、绘图

九、预测

PyTorch 常用代码段

1. 基本配置

导入包和版本查询

可复现性

显卡设置

2. 张量(Tensor)处理

张量的数据类型

张量基本信息

命名张量

数据类型转换

torch.Tensor与np.ndarray转换

Torch.tensor与PIL.Image转换

np.ndarray与PIL.Image的转换

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

张量形变

打乱顺序

水平翻转

复制张量

张量拼接

将整数标签转为one-hot编码

得到非零元素

判断两个张量相等

张量扩展

矩阵乘法

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

3. 模型定义和操作

一个简单两层卷积网络的示例

双线性汇合（bilinear pooling）

多卡同步 BN（Batch normalization）

将已有网络的所有BN层改为同步BN层

类似 BN 滑动平均

计算模型整体参数量

查看网络中的参数

模型可视化（使用pytorchviz）

类似 Keras 的 model.summary() 输出模型信息（使用pytorch-summary ）

提取模型中的某一层

部分层使用预训练模型

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

4. 数据处理

计算数据集的均值和标准差

得到视频数据基本信息

TSN 每段（segment）采样一帧视频

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

5. 模型训练和测试

分类模型训练代码

分类模型测试代码

自定义loss

标签平滑（label smoothing）

Mixup训练

L1 正则化

不对偏置项进行权重衰减（weight decay）

梯度裁剪（gradient clipping）

得到当前学习率

学习率衰减

优化器链式更新

模型训练可视化

保存与加载断点

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

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

微调全连接层

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

6. 其他注意事项

PyTorch使用技巧

目录

指定GPU编号

查看模型每层输出详情

梯度裁剪（Gradient Clipping）

扩展单张图片维度

独热编码