python tqdm pytorch_pytorch使用horovod多gpu训练的实现

pytorch在Horovod上训练步骤分为以下几步：

import torch

import horovod.torch as hvd

# Initialize Horovod 初始化horovod

hvd.init()

# Pin GPU to be used to process local rank (one GPU per process) 分配到每个gpu上

torch.cuda.set_device(hvd.local_rank())

# Define dataset... 定义dataset

train_dataset = ...

# Partition dataset among workers using DistributedSampler 对dataset的采样器进行调整，使用torch.utils.data.distributed.DistributedSampler

train_sampler = torch.utils.data.distributed.DistributedSampler(

train_dataset, num_replicas=hvd.size(), rank=hvd.rank())

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

# Build model...

model = ...

model.cuda()

optimizer = optim.SGD(model.parameters())

# Add Horovod Distributed Optimizer 使用Horovod的分布式优化器函数包裹在原先optimizer上

optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())

# Broadcast parameters from rank 0 to all other processes. 参数广播到每个gpu上

hvd.broadcast_parameters(model.state_dict(), root_rank=0)

for epoch in range(100):

for batch_idx, (data, target) in enumerate(train_loader):

optimizer.zero_grad()

output = model(data)

loss = F.nll_loss(output, target)

loss.backward()

optimizer.step()

if batch_idx % args.log_interval == 0:

print('Train Epoch: {} [{}/{}]\tLoss: {}'.format(

epoch, batch_idx * len(data), len(train_sampler), loss.item()))

完整示例代码如下，在imagenet上采用resnet50进行训练

from __future__ import print_function

import torch

import argparse

import torch.backends.cudnn as cudnn

import torch.nn.functional as F

import torch.optim as optim

import torch.utils.data.distributed

from torchvision import datasets, transforms, models

import horovod.torch as hvd

import os

import math

from tqdm import tqdm

from distutils.version import LooseVersion

# Training settings

parser = argparse.ArgumentParser(description='PyTorch ImageNet Example',

formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('--train-dir', default=os.path.expanduser('~/imagenet/train'),

help='path to training data')

parser.add_argument('--val-dir', default=os.path.expanduser('~/imagenet/validation'),

help='path to validation data')

parser.add_argument('--log-dir', default='./logs',

help='tensorboard log directory')

parser.add_argument('--checkpoint-format', default='./checkpoint-