只保存参数信息
加载
checkpoint = torch.load(opt.resume)
model.load_state_dict(checkpoint)
保存
torch.save(self.state_dict(),file_path)
这而只保存了参数信息,读取时也只有参数信息,模型结构需要手动编写
保存整个模型
保存
torch.save(the_model, PATH)
加载:
the_model = torch.load(PATH)
有时候会看到加载时
model.load_state_dict(checkpoint['state_dic'])
这是因为checkpoint是一个字典,保存的key可以自己定义。
可以保存除参数信息之外的其它信息,如epoch等。
保存:
torch.save({
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
}, 'checkpoint.tar' )
加载:
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.evaluate, checkpoint['epoch']))
model.state_dict就是一个字典,可以遍历
for k,v in model.state_dict():
print(k)
print(v)