文章目錄
一、訓練完整流程
使用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加載時會報錯。