Pytorch--分布式训练(下)

最新推荐文章于 2025-06-15 11:27:09 发布

置顶

Never-guess

最新推荐文章于 2025-06-15 11:27:09 发布

阅读量1.8w

点赞数 1

CC 4.0 BY-SA版权

分类专栏：深度学习文章标签： pytorch 分布式深度学习教程代码

本文链接：https://blog.youkuaiyun.com/qq_20791919/article/details/79057871

看过 Pytorch–分布式训练(上)的朋友肯定有一个疑问，怎么才能实现多节点多gpu中的分布式训练呢？
同样是在官网，找到这样一个imagenet分布式训练的例子
https://github.com/pytorch/examples/tree/master/imagenet
然后可能是由于pytorch版本更新问题，我一开始并没有跑通程序，后来修改了一下就可以了。
修改版本
https://github.com/zzy123abc/distributed-pytorch

import argparse
import os
import shutil
import time

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models

model_names = sorted(name for name in models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(models.__dict__[name]))

parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('data', metavar='DIR',
                    help='path to dataset')
parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18',
                    choices=model_names,
                    help='model architecture: ' +
                        ' | '.join(model_names) +
                        ' (default: resnet18)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--epochs', default=90, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=256, type=int,
                    metavar='N', help='mini-batch size (default: 256)')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq', '-p', default=10, type=int,
                    metavar=