如何以正常的非并行方式运行 Pytorch 模型？

Question

我正在查看this script，这里有一个代码块考虑了两个选项，DataParallel 和DistributedDataParallel：

if not args.distributed:
    if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
        model.features = torch.nn.DataParallel(model.features)
        model.cuda()
    else:
        model = torch.nn.DataParallel(model).cuda()
else:
    model.cuda()
    model = torch.nn.parallel.DistributedDataParallel(model)

如果我不想要这些选项中的任何一个，并且我想在没有DataParallel 的情况下运行它。我该怎么做？

如何定义我的模型，以便它以普通的nn 运行而不并行化任何东西？

Answer 1

• DataParallel 是一个包装对象，用于在同一台机器的多个 GPU 上并行计算，请参阅 here。
• DistributedDataParallel 也是一个包装对象，可让您在多个设备上分发数据，请参阅 here。

如果您不想要它，您可以简单地移除包装器并按原样使用模型：

if not args.distributed:
    if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
        model.features = model.features
        model.cuda()
    else:
        model = model.cuda()
else:
    model.cuda()
    model = model

这是为了尽量减少代码修改。当然，由于您对并行化不感兴趣，因此您可以将整个 if 语句删除为以下内容：

if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
    model.features = model.features
model = model.cuda()

请注意，此代码假定您在 GPU 上运行。