pytorch实现单机多卡训练

pytorch实现单机多卡训练

from torch.nn.parallel import DataParallel # 单机多卡的分布式训练（数据并行） 模型训练加速

    '''
    三 构建模型
    '''
    model = build_model(CFG.backbone, CFG.num_classes, CFG.device)
    model.load_state_dict(torch.load("best_epoch.bin"))
    model = DataParallel(model)