model = model.to(device)
if args.resume:
print("恢复训练")
print("============================================>")
checkpoint = torch.load(args.resume,map_location=torch.device('cpu')) #不加cpu会占用 gpu 内存, gpu 快满了,会导致 outof memry
start_epoch = checkpoint['epoch'] + 1
model.load_state_dict(checkpoint['model_state_dict']) # 加载模型可学习参数
optimizer.load_state_dict(checkpoint['optimizer_state_dict']) #
pytorch加载预选连模型 或者中断恢复训练 gpu outof memory
最新推荐文章于 2024-05-16 14:59:34 发布