便纵有千种风情

更与何人说

PyTorc读取和保存模型参数

只保存参数信息

加载

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)
阅读更多
个人分类: PyTorch
想对作者说点什么? 我来说一句

没有更多推荐了,返回首页

加入CSDN,享受更精准的内容推荐,与500万程序员共同成长!
关闭
关闭