pytorch保存模型並記錄最優模型

 

# https://github.com/tczhangzhi/pytorch-distributed/blob/master/distributed.py

# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)

if args.local_rank == 0:
    save_checkpoint(
               {
                    'epoch': epoch + 1,
                    'arch': args.arch,
                    'state_dict': model.module.state_dict(),
                    'best_acc1': best_acc1,
                }, is_best)



def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pth.tar')

　　

shutil.copyfile(filename, 'model_best.pth.tar') # 如果是當前最優精度的模型，則保存時維護一個副本