文章目录
一、训练完整流程
使用Pytorch训练神经网络的一般流程为（伪代码，许多功能需要自己实现，这里只列出了流程）：
import torch
import torch.optim as optim
import numpy as np
from tqdm import tqdm
from tensorboardX import SummaryWriter
from torch.optim.lr_scheduler import StepLR
def train():
with torch.cuda.device(gpu_id):
## model
model = Model() # 定义模型
model = model.cuda()
# optimizer & lr_scheduler
optimizer = torch.optim.SGD(model.parameters(), lr=init_lr,
momentum=momentum, weight_decay=weight_decay)
lr_scheduler = StepLR(optimizer, step_size=25, gamma=0.8) # 定义学习率
# lr_scheduler = lr_decay() # 也可以是自己定义的学习率下降方式,比如定义了一个列表
if resume: # restore from checkpoint
model, optimizer = restore_from(model, optimizer, ckpt) # 恢复训练状态
# load train data
trainloader, validloader = dataset() #自己定义DataLoader
### logs
logger = create_logger() # 自己定义创建的log日志
summary_writer = SummaryWriter(log_dir) # tensorboard
### start train
for epoch in range(end_epoch):
scheduler.step() # 更新optimizer的学习率，一般以epoch为单位，即多少个epoch后换一次学习率
train_loss = []
model.train()
model = model.cuda()
## train
for i, data in enumerate(tqdm(trainloader)):
input, target = data
optimizer.zero_grad() #使用之前先清零
output = model(input.cuda())
loss = Loss(output, target) # 自己定义损失函数
loss.backward() # loss反传，计算模型中各tensor的梯度
optimizer.step() #用在每一个mini batch中，只有用了optimizer.step()，模型才会更新
train_loss.append(loss)
train_loss = np.mean(train_loss) # 对各个mini batch的loss求平均
## eval，不需要梯度反传
valid_loss = []
model.eval() # 注意model的模式从train()变成了eval()
for i, data in enumerate(tqdm(validloader)):
input, target = data
optimizer.zero_grad()
output = model(input.cuda())
loss = Loss(output, target) # 自己定义损失函数
valid_loss.append(loss)
valid_loss = np.mean(valid_loss)
summary_writer.add_scalars('loss', {'train_loss': train_loss, 'valid_loss': valid_loss}, epoch) #写入tensorboard
if (epoch + 1) % 10 == 0 or (epoch + 1) == end_epoch: # 保存模型
torch.save(
{'epoch': epoch,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict()},
save_path)
def restore_from(model, optimizer, ckpt_path):
device = torch.cuda.current_device()
ckpt = torch.load(ckpt_path, map_location=lambda storage, loc: storage.cuda(device))
epoch = ckpt['epoch']
ckpt_model_dict = remove_prefix(ckpt['state_dict'], 'module.')
model.load_state_dict(ckpt_model_dict, strict=False) # load model
optimizer.load_state_dict(ckpt['optimizer']) # load optimizer
return model, optimizer, epoch
def remove_prefix(state_dict, prefix):
''' Old style model is stored with all names of parameters share common prefix 'module.' '''
logger.info('remove prefix \'{}\''.format(prefix))
f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
return {f(key): value for key, value in state_dict.items()}
二、高阶操作
1.自定义学习率
用torch.optim.lr_scheduler里面自带的函数设定学习率（例如StepLR），返回的lr_schduler是一个对象<torch.optim.lr_scheduler.StepLR object at 0x000001C3CE3A12E8>，可以使用scheduler.step()来更新optimizer的学习率。如果我们要自定义学习率的下降方式，就不再具备scheduler.step()功能了，需要手动更新optimizer的学习率，举个栗子：
这里在[init_lr,end_lr]范围内，定义了log状下降的学习率,返回的lr_scheduler是一个长度为end_epoch的numpy array
import math
init_lr = 0.01
end_lr = 0.0001
end_epoch = 10
lr_scheduler = np.logspace(math.log10(init_lr), math.log10(end_lr),end_epoch)
# lr_scheduler ：[0.01 0.00599484 0.00359381 0.00215443 0.00129155 \
# 0.00077426 0.00046416 0.00027826 0.00016681 0.0001 ]
将学习率应用于每轮训练中，用于更新optimizer学习率的方法：在每个epoch开始时候，对optimizer参数中的lr参数进行手动调整。
for epoch in range(end_epoch):
curLR = lr_scheduler[epoch] # 用于本轮训练的lr
for param_group in optimizer.param_groups:
param_group['lr'] = curLR
2. 只训练特定的网络层
我们来看看上面定义的optimizer：
optimizer = torch.optim.SGD(model.parameters(), lr=init_lr,
momentum=momentum, weight_decay=weight_decay)
这是将model的所有参数都用来训练，其中model.parameters()代表的是网络的所有参数。如果只需要训练特定的几层（特定的参数），应该改成：
trainable_params = [p for p in model.parameters() if p.requires_grad] # 获取所有可训练参数（可以求梯度的参数是可训练参数）
optimizer = torch.optim.SGD(trainable_params, lr=init_lr,
momentum=momentum, weight_decay=weight_decay)
这样即可对自己定义的某些参数进行训练（想要训练哪些参数，将这些参数的requires_grad改成True。）
3. 逐层释放/冻结网络参数
这种操作也是很常见的，训练过程中可能需要训练更多/更少的参数，因为可训练的参数量发生了变化，原来的optimizer不能再继续用了。所以再释放/冻结可训练参数的同时，需要重新定义optimizer（这里以释放参数举例）：
class Model:
def unfix(self, epoch): # 随手写了个示例，例如当epoch=5和10时释放新参数，epoch为其他时不释放
if epoch == 5:
# 释放一部分参数,将新加入的训练参数的requires_grad设成True
return True
elif epoch == 10:
# 释放一部分参数,将新加入的训练参数的requires_grad设成True
return True
else:
return False
if model.unfix(epoch): # 如果当前epoch训练参数量发生了变化，需要重新定义optimizer
trainable_params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(trainable_params, lr=init_lr,
momentum=momentum, weight_decay=weight_decay)
4.恢复优化器状态时参数不match的解决方案
最开始的训练流程里面提到过，使用torch.load(ckpt)可以恢复模型，再使用model.load_state_dict(ckpt['state_dict'])和optimizer.load_state_dict(ckpt['optimizer'])就可以分别恢复model和optimizer的状态了。但如果在训练过程中，我们使用了逐层释放/冻结网络参数的训练方式，会对应有多个optimizer，当前的optimizer可能会跟恢复的optimizer不一致，报错为："loaded state dict contains a parameter group " "that doesn't match the size of optimizer's group"
例如，我们想从epoch轮开始继续训练，需要resotre epoch-1轮保存好的模型。而如果第epoch轮正好有参数释放，就会导致epoch和epoch-1对应的训练参数量不相同，解决方法是：
（1）找到待恢复的模型中optimizer的训练参数，先把optimizer恢复出来
（2）然后再根据当前epoch，判断是否需要释放新的参数
def get_trainable_params(model, epoch):
flag = model.unfix((epoch - 1)) # 看上一轮有没有释放参数
if flag:
trainable_params = [p for p in model.parameters() if p.requires_grad]
return trainable_params
else: # 这轮没有释放，找到最近的上次释放的地方
for i in range(epoch - 1, 0, -1):
flag = model.unfix(i)
if flag:
trainable_params = [p for p in model.parameters() if p.requires_grad]
return trainable_params
else:
continue
return
trainable_params = get_trainable_params(model, epoch) # 找到待恢复的模型中optimizer的训练参数
optimizer = torch.optim.SGD(trainable_params, lr=init_lr,
momentum=momentum, weight_decay=weight_decay)
model, optimizer = restore_from(model, optimizer, ckpt) # 先恢复之前的旧模型
if model.unfix(epoch): # 如果当前epoch训练参数量发生了变化，需要重新定义optimizer
trainable_params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(trainable_params, lr=init_lr,
momentum=momentum, weight_decay=weight_decay)
5. 梯度反传，loss反传，梯度裁剪
前面的训练流程里，在计算完loss之后，直接对loss进行了反传和梯度的更新，这样可能会出现一个问题：如果loss值为nan，或者梯度消失或爆炸，会对网络训练产生很大的影响。因此可以这样修改代码：
loss = Loss(output, target) # 自己定义损失函数
if is_valid_number(loss.data.item()): # 判断loss是否合法
loss.backward()
# clip gradient
clip_grad_norm_(model.parameters(), cfg.TRAIN.GRAD_CLIP)
optimizer.step()
def is_valid_number(x):
return not(math.isnan(x) or math.isinf(x) or x > 1e4)
三、恢复保存的优化器状态，继续优化
#### 保存模型
states_dict = {'epoch': epoch, 'arch': model_name, 'model':model.module.state_dict(),'optimizer': optimizer.state_dict()}
torch.save(states_dict, 'model.pth') # save model
#### 加载/恢复模型
ckpt = torch.load('model.pth', map_location=lambda storage, loc: storage) # 读取/加载模型
epoch = ckpt['epoch']
arch = ckpt['arch']
model.load_state_dict(ckpt['model']) # load model
optimizer.load_state_dict(ckpt['optimizer']) # load optimizer
# We must convert the resumed state data of optimizer to gpu
"""It is because the previous training was done on gpu, so when saving the optimizer.state_dict, the stored
states(tensors) are of cuda version. During resuming, when we load the saved optimizer, load_state_dict()
loads this cuda version to cpu. But in this project, we use map_location to map the state tensors to cpu.
In the training process, we need cuda version of state tensors, so we have to convert them to gpu."""
for state in optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.cuda()
总结：在训练的过程中，使用torch.save()存的是cuda型，恢复训练的时候，使用load_state_dict()函数，会将cuda型转成cpu型，因此为了恢复训练，需要将这些数据再转成cuda型。
四、加载模型到指定的卡上
自己在load模型，继续训练的时候遇到了这样的问题：
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)
网上查资料发现，一个博主分析说：pytorch模型中会记录GPU信息，所以如果测试时候使用不同于训练的gpu加载模型，则会报错。
注意一个细节：训练模型时使用的GPU卡和加载时使用的GPU卡不一样导致的。个人感觉，因为pytorch的模型中是会记录有GPU信息的，所以有时使用不同的GPU加载时会报错。