轉自：知乎

目錄：

• 儲存模型與載入模型
• 凍結一部分引數，訓練另一部分引數
• 採用不同的學習率進行訓練

1.儲存模型與載入

簡單的儲存與載入方法：

# 儲存整個網路

torch.save(net, PATH)

# 儲存網路中的引數, 速度快，佔空間少

torch.save(net.state_dict(),PATH)

#--------------------------------------------------

#針對上面一般的儲存方法，載入的方法分別是：

model_dict=torch.load(PATH)

model_dict=model.load_state_dict
 
(torch.load(PATH))

然而，在實驗中往往需要儲存更多的資訊，比如優化器的引數，那麼可以採取下面的方法儲存：

torch.save({'epoch': epochID + 1, 'state_dict': model.state_dict(), 'best_loss': lossMIN,

'optimizer': optimizer.state_dict(),'alpha': loss.alpha, 'gamma': loss.gamma},

checkpoint_path + '/m-' + launchTimestamp + '-' + str("%.4f" % lossMIN) + '.pth.tar')

 

以上包含的資訊有，epochID, state_dict, min loss, optimizer, 自定義損失函式的兩個引數；格式以字典的格式儲存。

載入的方式：

def load_checkpoint(model, checkpoint_PATH, optimizer):

if checkpoint != None:

model_CKPT = torch.load(checkpoint_PATH)

model.load_state_dict(model_CKPT['state_dict'])

print('loading checkpoint!')

optimizer.load_state_dict(model_CKPT['optimizer'])

return model, optimizer

 

其他的引數可以通過以字典的方式獲得

但是，但是，我們可能修改了一部分網路，比如加了一些，刪除一些，等等，那麼需要過濾這些引數，載入方式：

def load_checkpoint(model, checkpoint, optimizer, loadOptimizer):

if checkpoint != 'No':

print("loading checkpoint...")

model_dict = model.state_dict()

modelCheckpoint = torch.load(checkpoint)

pretrained_dict = modelCheckpoint['state_dict']

# 過濾操作

new_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict.keys()}

model_dict.update(new_dict)

# 打印出來，更新了多少的引數

print('Total : {}, update: {}'.format(len(pretrained_dict), len(new_dict)))

model.load_state_dict(model_dict)

print("loaded finished!")

# 如果不需要更新優化器那麼設定為false

if loadOptimizer == True:

optimizer.load_state_dict(modelCheckpoint['optimizer'])

print('loaded! optimizer')

else:

print('not loaded optimizer')

else:

print('No checkpoint is included')

return model, optimizer

2.凍結部分引數，訓練另一部分引數

1）新增下面一句話到模型中

for p in self.parameters():

p.requires_grad = False

比如載入了resnet預訓練模型之後，在resenet的基礎上連線了新的模快，resenet模組那部分可以先暫時凍結不更新，只更新其他部分的引數，那麼可以在下面加入上面那句話

class RESNET_MF(nn.Module):

def init(self, model, pretrained):

super(RESNET_MF, self).__init__()

self.resnet = model(pretrained)

for p in self.parameters():

p.requires_grad = False

self.f = SpectralNorm(nn.Conv2d(2048, 512, 1))

self.g = SpectralNorm(nn.Conv2d(2048, 512, 1))

self.h = SpectralNorm(nn.Conv2d(2048, 2048, 1))

...

同時在優化器中新增：filter(lambda p: p.requires_grad, model.parameters())

optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001, betas=(0.9, 0.999),

eps=1e-08, weight_decay=1e-5)

2) 引數儲存在有序的字典中，那麼可以通過查詢引數的名字對應的id值，進行凍結

查詢的程式碼：

 model_dict = torch.load('net.pth.tar').state_dict()

dict_name = list(model_dict)

for i, p in enumerate(dict_name):

print(i, p)

儲存一下這個檔案，可以看到大致是這個樣子的：

0 gamma

1 resnet.conv1.weight

2 resnet.bn1.weight

3 resnet.bn1.bias

4 resnet.bn1.running_mean

5 resnet.bn1.running_var

6 resnet.layer1.0.conv1.weight

7 resnet.layer1.0.bn1.weight

8 resnet.layer1.0.bn1.bias

9 resnet.layer1.0.bn1.running_mean

....

同樣在模型中新增這樣的程式碼：

for i,p in enumerate(net.parameters()):

if i < 165:

p.requires_grad = False

在優化器中新增上面的那句話可以實現引數的遮蔽