一、保存神经网络模型
Pytorch模型以pth格式进行保存,为了方便操作我们可以设置两种保存模式:保存一定轮次的模型数据和保存验证集损失函数最小的模型数据。同时也可保存损失函数日志以便分析。代码实现如下:
torch.save(model.state_dict(), os.path.join(保存路径, "文件名"))
损失函数可以保存在txt文件中
def append_loss(self, epoch, loss, val_loss):
if not os.path.exists(self.log_dir):
os.makedirs(self.log_dir)
self.losses.append(loss)
self.val_loss.append(val_loss)
with open(os.path.join(self.log_dir, "epoch_loss.txt"), 'a') as f:
f.write(str(loss))
f.write("\n")
with open(os.path.join(self.log_dir, "epoch_val_loss.txt"), 'a') as f:
f.write(str(val_loss))
f.write("\n")
self.writer.add_scalar('loss', loss, epoch)
self.writer.add_scalar('val_loss', val_loss, epoch)
#可以添加图像绘制代码
轮次保存和最佳保存可以添加逻辑代码予以控制。
if (epoch + 1) % save_period == 0 or epoch + 1 == Epoch:
#保存代码
if len(loss_history.val_loss) <= 1 or (val_loss / epoch_step_val) <= min(loss_history.val_loss):
#保存代码
二、神经网络模型加载
加载的模型一般有两类,一种是直接由pytorch保存的.pth文件,另一种是通用文件.onnx。其中onnx是微软推出的通用网络模型标准,可以实现跨平台部署。onnx文件的加载会在以后描述。
#直接加载
self.net.load_state_dict(torch.load('文件路径', map_location=device))
#分层加载--可以检测不匹配的层
model_dict = model.state_dict() #加载模型字典用于比对
pretrained_dict = torch.load(model_path, map_location = device)
load_key, no_load_key, temp_dict = [], [], {}
for k, v in pretrained_dict.items():
if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
temp_dict[k] = v
load_key.append(k) #匹配成功的层
else:
no_load_key.append(k) #匹配失败的层
model_dict.update(temp_dict)
model.load_state_dict(model_dict)