pytorch加载预选连模型 或者中断恢复训练 gpu outof memory

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'])  #

加 map_location=torch.device(‘cpu’) ,因为正好gpu 快满了, 加载参数默认放在gpu上占用gpu参数, 所以显卡爆了, 这里加载的参数放到 cpu 上就行了,

你可能感兴趣的:(pytorch加载预选连模型 或者中断恢复训练 gpu outof memory)