pytorch梯度调试以及模型的保存和加载

1、打印梯度

for name, param in model.named_parameters():
    if param.requires_grad:
        if param.grad is not None:
            print("{}, gradient: {}".format(name, param.grad.mean()))
        else:
            print("{} has not gradient".format(name))

 

2、保存模型和加载模型

# 保存模型到 filename
torch.save({
'epoch': epoch + 1, 'model': model.module.state_dict() if args.mGPUs else model.state_dict(), 'optimizer': optimizer.state_dict(), }, filename)
# 从 filename 加载模型
checkpoint = torch.load(filename)
start_epoch = checkpoint['epoch']
checkpoint['model'] = {k: v for k, v in chkpt['model'].items() if model.state_dict()[k].numel() == v.numel()}
model.load_state_dict(checkpoint['model'], strict=False)
optimizer.load_state_dict(checkpoint['optimizer'])

 

posted @ 2020-06-04 16:24  blackx  阅读(2068)  评论(0编辑  收藏  举报