PyTorch | 模型加載/參數初始化/Finetune

1、模型保存/加載

1.1 全部模型參數

訓練過程當中，有時候會因爲各類緣由中止訓練，這時候咱們訓練過程當中就須要注意將每一輪epoch的模型保存(通常保存最好模型與當前輪模型)。通常使用pytorch裏面推薦的保存方法。該方法保存的是模型的參數。

#保存模型到checkpoint.pth.tar
torch.save(model.module.state_dict(), ‘checkpoint.pth.tar’)

對應的加載模型方法爲(這種方法須要先反序列化模型獲取參數字典，所以必須先load模型，再load_state_dict)：

mymodel.load_state_dict(torch.load(‘checkpoint.pth.tar’))

有了上面的保存後，現以一個例子說明如何在inference AND/OR resume train使用。

#保存模型的狀態，能夠設置一些參數，後續可使用
state = {'epoch': epoch + 1,#保存的當前輪數
         'state_dict': mymodel.state_dict(),#訓練好的參數
         'optimizer': optimizer.state_dict(),#優化器參數,爲了後續的resume
         'best_pred': best_pred#當前最好的精度
          ,....,...}

#保存模型到checkpoint.pth.tar
torch.save(state, ‘checkpoint.pth.tar’)
#若是是best,則複製過去
if is_best:
    shutil.copyfile(filename, directory + 'model_best.pth.tar')

checkpoint = torch.load('model_best.pth.tar')
model.load_state_dict(checkpoint['state_dict'])#模型參數
optimizer.load_state_dict(checkpoint['optimizer'])#優化參數
epoch = checkpoint['epoch']#epoch，能夠用於更新學習率等

#有了以上的東西，就能夠繼續從新訓練了，也就不須要擔憂中止程序從新訓練。
train/eval
....
....

1.2 部分模型參數

在不少時候，咱們加載的是已經訓練好的模型，而訓練好的模型可能與咱們定義的模型不徹底同樣，而咱們只想使用同樣的那些層的參數。

有幾種解決方法：

（1）直接在訓練好的模型開始搭建本身的模型，就是先加載訓練好的模型，而後再它基礎上定義本身的模型；

model_ft = models.resnet18(pretrained=use_pretrained)
self.conv1 = model_ft.conv1
self.bn = model_ft.bn
... ...

（2) 本身定義好模型，直接加載模型

#第一種方法：
mymodelB = TheModelBClass(*args, **kwargs)
# strict=False，設置爲false,只保留鍵值相同的參數
mymodelB.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)

#第二種方法：
# 加載模型
model_pretrained = models.resnet18(pretrained=use_pretrained)

# mymodel's state_dict，
# 如：  conv1.weight 
#     conv1.bias  
mymodelB_dict = mymodelB.state_dict()

# 將model_pretrained的建與自定義模型的建進行比較，剔除不一樣的
pretrained_dict = {k: v for k, v in model_pretrained.items() if k in mymodelB_dict}
# 更新現有的model_dict
mymodelB_dict.update(pretrained_dict)

# 加載咱們真正須要的state_dict
mymodelB.load_state_dict(mymodelB_dict)

# 方法2可能更直觀一些

2、參數初始化

第二個問題是參數初始化問題，在不少代碼裏面都會使用到，畢竟不是全部的都是有預訓練參數。這時就須要對不是與預訓練參數進行初始化。pytorch裏面的每一個Tensor實際上是對Variabl的封裝，其包含data、grad等接口，所以能夠用這些接口直接賦值。這裏也提供了怎樣把其餘框架(caffe/tensorflow/mxnet/gluonCV等)訓練好的模型參數直接賦值給pytorch.其實就是對data直接賦值。

pytorch提供了初始化參數的方法：

 def weight_init(m):
    if isinstance(m,nn.Conv2d):
        n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        m.weight.data.normal_(0,math.sqrt(2./n))
    elif isinstance(m,nn.BatchNorm2d):
        m.weight.data.fill_(1)
        m.bias.data.zero_()

但通常若是沒有很大需求初始化參數，也沒有問題(不肯定性能是否有影響的狀況下)，pytorch內部是有默認初始化參數的。

3、模型微調/Finetune

最後是微調，平時作實驗，至少backbone是用預訓練的模型，將其用做特徵提取器，或者在它上面作微調。

用於特徵提取的時候，要求特徵提取部分參數不進行學習，而pytorch提供了requires_grad參數用於肯定是否進去梯度計算，也便是否更新參數。如下以minist爲例，用resnet18做特徵提取：

#加載預訓練模型
model = torchvision.models.resnet18(pretrained=True)

#遍歷每個參數，將其設置爲不更新參數，即不學習
for param in model.parameters():
    param.requires_grad = False

# 將全鏈接層改成mnist所需的10類，注意：這樣更改後requires_grad默認爲True
model.fc = nn.Linear(512, 10)

# 優化
optimizer = optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9)         

用於全局微調時，咱們通常對不一樣的層須要設置不一樣的學習率，預訓練的層學習率小一點，其餘層大一點。這要怎麼作呢？

# 加載預訓練模型
model = torchvision.models.resnet18(pretrained=True)
model.fc = nn.Linear(512, 10)

# 參考：https://blog.csdn.net/u012759136/article/details/65634477
ignored_params = list(map(id, model.fc.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())

# 對不一樣參數設置不一樣的學習率
params_list = [{'params': base_params, 'lr': 0.001},]
params_list.append({'params': model.fc.parameters(), 'lr': 0.01})

optimizer = torch.optim.SGD(params_list,
                    0.001，
                    momentum=args.momentum,
                    weight_decay=args.weight_decay)

 

參考：

https://zhuanlan.zhihu.com/p/48524007

https://zhuanlan.zhihu.com/p/38056115