文章详情

短信预约-IT技能 免费直播动态提醒

请输入下面的图形验证码

提交验证

短信预约提醒成功

pytorch 多GPU并行训练代码讲解

2023-09-01 21:18

关注

文章目录

1. 介绍

pytorch 多GPU并行训练原理介绍,参考我之前的博客: pytorch中使用多GPU并行训练

1.1 多GPU训练的命令

(1)指定使用GPU:0运行脚本
方式1:

CUDA_VISIBLE_DEVICES=0 python ***.py

方式2:在python脚本文件中添加如下内容进行指定

import osos.environ["CUDA_VISIBLE_DEVICES"] = "0"

如果在脚本文件中已经写明了第二种方式,则不支持第一种方式的指定

(2) 指定使用多张显卡运行脚本
在GPU的id为0和1的两张显卡上运行***.py程序

CUDA_VISIBLE_DEVICES=0,1 python ***.py

(3) 指定所有显卡一起运行脚本

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --use_env train.py

如果提示错误:RuntimeError: The server socket has failed to listen on any local network address. The server socket has failed to bind to [::]:29500 (errno: 98 - Address already in use). The server socket has failed to bind to ?UNKNOWN? (errno: 98 - Address already in use).

则执行:

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --master_port='29501' --use_env train.py

(4) 新版的多GPU运行指令

torchrun --nproc_per_node=8 train.py
CUDA_VISIBLE_DEVICES=0,3 torchrun --nproc_per_node=2 train.py

1.2 查看GPU状态

nvidia-smi
watch -n 1 nvidia-smi

pytorch多GPU训练启动方式

第一种:torch.distributed.launch,相比于第二种torch.multiprocessing代码量会更少一点,启动速度会更快一点。

项目代码

https://github.com/WZMIAOMIAO/deep-learning-for-image-processing/tree/master/pytorch_classification/train_multi_GPU,项目中实现了两种GPU启动的方式:

2. 代码讲解

项目以ResNet网络花朵数据集分类为例说明。

2.1单GPU训练

train_multi_GPU中作为参考,首先提供了一个train_single_GPU.py使用单GPU训练的脚本。

训练参数说明

 parser = argparse.ArgumentParser()    parser.add_argument('--num_classes', type=int, default=5)    parser.add_argument('--epochs', type=int, default=30)    parser.add_argument('--batch-size', type=int, default=16)    parser.add_argument('--lr', type=float, default=0.001)    parser.add_argument('--lrf', type=float, default=0.1)    # 数据集所在根目录    # https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz    parser.add_argument('--data-path', type=str,                        default="/home/w180662/my_project/my_github/data_set/flower_data/flower_photos")    # resnet34 官方权重下载地址    # https://download.pytorch.org/models/resnet34-333f7ec4.pth    parser.add_argument('--weights', type=str, default='resNet34.pth',                        help='initial weights path')    parser.add_argument('--freeze-layers', type=bool, default=False)    parser.add_argument('--device', default='cuda', help='device id (i.e. 0 or 0,1 or cpu)')    opt = parser.parse_args()

训练脚本

main函数

 device = torch.device(args.device if torch.cuda.is_available() else "cpu")
print('Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/')tb_writer = SummaryWriter()
train_info, val_info, num_classes = read_split_data(args.data_path) train_images_path, train_images_label = train_info val_images_path, val_images_label = val_info # check num_classes assert args.num_classes == num_classes, "dataset num_classes: {}, input {}".format(args.num_classes,num_classes) data_transform = {     "train": transforms.Compose([transforms.RandomResizedCrop(224),      transforms.RandomHorizontalFlip(),      transforms.ToTensor(),      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),     "val": transforms.Compose([transforms.Resize(256),    transforms.CenterCrop(224),    transforms.ToTensor(),    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])} # 实例化训练数据集 train_data_set = MyDataSet(images_path=train_images_path,images_class=train_images_label,transform=data_transform["train"]) # 实例化验证数据集 val_data_set = MyDataSet(images_path=val_images_path,                          images_class=val_images_label,                          transform=data_transform["val"]) batch_size = args.batch_size nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers print('Using {} dataloader workers every process'.format(nw)) train_loader = torch.utils.data.DataLoader(train_data_set,                batch_size=batch_size,                shuffle=True,                pin_memory=True,                num_workers=nw,                collate_fn=train_data_set.collate_fn) val_loader = torch.utils.data.DataLoader(val_data_set,              batch_size=batch_size,              shuffle=False,              pin_memory=True,              num_workers=nw,              collate_fn=val_data_set.collate_fn)

1 ) 实例化模型,传入指定的类别数args.num_classes

 model = resnet34(num_classes=args.num_classes).to(device)

2 ) 载入模型权重:通过toch.load载入之后,它是有序字典的形式。 遍历权重的有序字典,每遍历一层都会与我们当前model它的权重信息进行对比,我们这里简单粗暴对比两者的权重个数是否一样的。如果一样的话就保存到load_weights_dict当中,如果不一样就不会保存到字段道中。原预训练权重是基于imagenet上训练的,所以全连接层类别个数是1000,我们这里是5,所以最后一层是不匹配的,因此不会将全连接层的权重保存到load_weights_dict中,这样载入就不会出问题。然后我们的model通过load_state_dict去载入权重字典。

if args.weights != "":      if os.path.exists(args.weights):          weights_dict = torch.load(args.weights, map_location=device)          load_weights_dict = {k: v for k, v in weights_dict.items()   if model.state_dict()[k].numel() == v.numel()}          print(model.load_state_dict(load_weights_dict, strict=False))      else:          raise FileNotFoundError("not found weights file: {}".format(args.weights))

3 ) 冻结权重:如果args.freeze_layers为True的话,就会冻结除fc(全连接)以外的整个网络,如果为False的话就直接训练整个网络权重。
4 ) 将需要梯度更新的参数(require_grad为True),对于参数require_grad为False说明已经被冻结了,就不需要去管它了。然后定义优化器optimizer,传入需要更新的参数,指定初始学习率lr,动量momentum,还有正则项weight_decay.接下来定义了cosine学习率的曲线,然后通过pytorch的LambdaLR方法,传入优化器,然后再指定一个学习率变化的函数,也就是我们定义的lambda函数

pg = [p for p in model.parameters() if p.requires_grad]optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=0.005)# Scheduler https://arxiv.org/pdf/1812.01187.pdflf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf  # cosinescheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

cosine学习率变化曲线
在这里插入图片描述
x*math.pi/args.epochs的取值在[0,pi]之间,因此math.cos(x*math.pi/args.epochs)取值范围[-1,1],对它加上1之后 1+math.cos(x*math.pi/args.epochs)的取值在[0,2] ,然后又除以2,取值范围[0,1],然后再乘以(1-args.lrf)加上args.lrf,最终的取值范围就在[args.lrf,1],假设初始化的学习率为为1,最终的学习率为初始化的args.lrf。我们这里设置args.lrf为0.1,也就是学习率会呈现余弦变换,最终迭代了epochs后,最终的学习率变为初始学习率的0.1(args.lrf)倍

def train_one_epoch(model, optimizer, data_loader, device, epoch):    model.train()    loss_function = torch.nn.CrossEntropyLoss()    mean_loss = torch.zeros(1).to(device)    optimizer.zero_grad()    # 在进程0中打印训练进度    if is_main_process():        data_loader = tqdm(data_loader, file=sys.stdout)    for step, data in enumerate(data_loader):        images, labels = data        pred = model(images.to(device))        loss = loss_function(pred, labels.to(device))        loss.backward()        loss = reduce_value(loss, average=True)        mean_loss = (mean_loss * step + loss.detach()) / (step + 1)  # update mean losses        # 在进程0中打印平均loss        if is_main_process():            data_loader.desc = "[epoch {}] mean loss {}".format(epoch, round(mean_loss.item(), 3))        if not torch.isfinite(loss):            print('WARNING: non-finite loss, ending training ', loss)            sys.exit(1)        optimizer.step()        optimizer.zero_grad()    # 等待所有进程计算完毕    if device != torch.device("cpu"):        torch.cuda.synchronize(device)    return mean_loss.item()

train_one_epoch主要的代码讲解如下:

1 ) 通过model.train()方法进入训练模式,定义损失函数CrossEntropyLoss,清空优化器的梯度信息optimizer.zero_grad()

 model.train()    loss_function = torch.nn.CrossEntropyLoss()    mean_loss = torch.zeros(1).to(device)    optimizer.zero_grad()

2 ) 判断当前进程是否是主进程,也就是进程0。由于我们是单GPU训练脚本,当前进程肯定是一个主进程。使用tqdm库来封装data_loader,调用的时候会打印进度条,当使用多GPU的时候,我们通过以下的命令,使得只在主进程打印进度条,在其他进程中不会打印进度条。

 # 在进程0中打印训练进度    if is_main_process():        data_loader = tqdm(data_loader, file=sys.stdout)

3 ) 遍历数据,获取图像数据和标签,并将图像指定相应的GPU设备中,传入model正向传播得到预测输出pred,将输出pred与真实标签label计算损失,然后将损失进行反向传播,然后对历史的损失求平均。

for step, data in enumerate(data_loader):        images, labels = data        pred = model(images.to(device))        loss = loss_function(pred, labels.to(device))        loss.backward()        loss = reduce_value(loss, average=True) # 这句是针对多GPU场景使用的,在单GPU中这句话其实是不起作用的        mean_loss = (mean_loss * step + loss.detach()) / (step + 1)  # update mean losses # 更新平均损失

通过data_loader的desc参数,这样的话就可以显示当前平均损失的值了

 data_loader.desc = "[epoch {}] mean loss {}".format(epoch, round(mean_loss.item(), 3))

4 ) 通过optimizer.step(),利用optimizer优化器去对该step的参数进行更新,然后将优化器的梯度清空。

optimizer.step()optimizer.zero_grad()

同步多个设备之间的进度,最后会返回该轮的平均损失。

# 等待所有进程计算完毕if device != torch.device("cpu"):    torch.cuda.synchronize(device)

执行完train_one_epoch,返回train_single_gpu.py脚本。

执行 schedler.step(),根据定义好的schedler调整策略去更新学习率。然后进入evaluate方法中

def evaluate(model, data_loader, device):    model.eval()    # 用于存储预测正确的样本个数    sum_num = torch.zeros(1).to(device)    # 在进程0中打印验证进度    if is_main_process():        data_loader = tqdm(data_loader, file=sys.stdout)    for step, data in enumerate(data_loader):        images, labels = data        pred = model(images.to(device))        pred = torch.max(pred, dim=1)[1]        sum_num += torch.eq(pred, labels.to(device)).sum()    # 等待所有进程计算完毕    if device != torch.device("cpu"):        torch.cuda.synchronize(device)    sum_num = reduce_value(sum_num, average=False)    return sum_num.item()
# 在进程0中打印验证进度   if is_main_process():       data_loader = tqdm(data_loader, file=sys.stdout)   for step, data in enumerate(data_loader):       images, labels = data       pred = model(images.to(device))       pred = torch.max(pred, dim=1)[1]       sum_num += torch.eq(pred, labels.to(device)).sum()
# 等待所有进程计算完毕if device != torch.device("cpu"):    torch.cuda.synchronize(device)sum_num = reduce_value(sum_num, average=False)

执行完 evaluate函数之后,获得正确的样本数,然后除以总的样本个数,就得到准确率了。

acc = sum_num / (val_data_set)

然后通过tensorboard将当前epoch求得的mean_loss,还有准确率acc,还有学习率lr,将他们写入到tensorboard中,最后保存下当前epoch训练完后的权重

 tags = ["loss", "accuracy", "learning_rate"]        tb_writer.add_scalar(tags[0], mean_loss, epoch)        tb_writer.add_scalar(tags[1], acc, epoch)        tb_writer.add_scalar(tags[2], optimizer.param_groups[0]["lr"], epoch)        torch.save(model.state_dict(), "./weights/model-{}.pth".format(epoch))

2.2多GPU训练

对应的脚本为torch_multi_gpu_using_launch.py, 另外一个多GPU的脚本train_multi_gpu_using_spawn.py有兴趣可以看看,其实是差不多的

训练参数说明

 parser = argparse.ArgumentParser() parser.add_argument('--num_classes', type=int, default=5) parser.add_argument('--epochs', type=int, default=30) parser.add_argument('--batch-size', type=int, default=16) parser.add_argument('--lr', type=float, default=0.001) parser.add_argument('--lrf', type=float, default=0.1) # 是否启用SyncBatchNorm parser.add_argument('--syncBN', type=bool, default=True) # 数据集所在根目录 # https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz parser.add_argument('--data-path', type=str, default="/home/wz/data_set/flower_data/flower_photos") # resnet34 官方权重下载地址 # https://download.pytorch.org/models/resnet34-333f7ec4.pth parser.add_argument('--weights', type=str, default='resNet34.pth',                     help='initial weights path') parser.add_argument('--freeze-layers', type=bool, default=False) # 不要改该参数,系统会自动分配 parser.add_argument('--device', default='cuda', help='device id (i.e. 0 or 0,1 or cpu)') # 开启的进程数(注意不是线程),不用设置该参数,会根据nproc_per_node自动设置 parser.add_argument('--world-size', default=4, type=int,                     help='number of distributed processes') parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training') opt = parser.parse_args()

主函数介绍

if torch.cuda.is_available() is False:    raise EnvironmentError("not find GPU device for training.")
def init_distributed_mode(args):    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:        args.rank = int(os.environ["RANK"])        args.world_size = int(os.environ['WORLD_SIZE'])        args.gpu = int(os.environ['LOCAL_RANK'])    elif 'SLURM_PROCID' in os.environ:        args.rank = int(os.environ['SLURM_PROCID'])        args.gpu = args.rank % torch.cuda.device_count()    else:        print('Not using distributed mode')        args.distributed = False        return    args.distributed = True    torch.cuda.set_device(args.gpu)    args.dist_backend = 'nccl'  # 通信后端,nvidia GPU推荐使用NCCL    print('| distributed init (rank {}): {}'.format(        args.rank, args.dist_url), flush=True)    dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,world_size=args.world_size, rank=args.rank)    dist.barrier()

可以看到这函数中有RANK,WORLD_SIZE,因为distribute_dataparallel它是可以使用在多机多卡场景,也可以使用在单机多卡场景。

初始化进程环境后,对参数进行赋值,其中学习率args.lr *=args.world_size,当我们使用多GPU并行的时候,由于梯度一般在多块GPU上取一个均值,所以一般我们会对学习率进行增大。这里就简单粗暴的进行设置,用几块GPU就将学习率扩大几倍。

 args.lr *= args.world_size  # 学习率要根据并行GPU的数量进行倍增
if rank == 0:  # 在第一个进程中打印信息,并实例化tensorboard   print(args)   print('Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/')   tb_writer = SummaryWriter()   if os.path.exists("./weights") is False:       os.makedirs("./weights")

2 ) 数据预处理和实例化数据集

train_info, val_info, num_classes = read_split_data(args.data_path)train_images_path, train_images_label = train_infoval_images_path, val_images_label = val_info# check num_classesassert args.num_classes == num_classes, "dataset num_classes: {}, input {}".format(args.num_classes,                           num_classes)data_transform = {    "train": transforms.Compose([transforms.RandomResizedCrop(224),     transforms.RandomHorizontalFlip(),     transforms.ToTensor(),     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),    "val": transforms.Compose([transforms.Resize(256),   transforms.CenterCrop(224),   transforms.ToTensor(),   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}# 实例化训练数据集train_data_set = MyDataSet(images_path=train_images_path,                           images_class=train_images_label,                           transform=data_transform["train"])# 实例化验证数据集val_data_set = MyDataSet(images_path=val_images_path,                         images_class=val_images_label,                         transform=data_transform["val"])

这些代码是一样的,是单GPU不一样的,这里使用了DistributedSampler,这个函数通过鼠标左键找不到它的定义。我们可以复制DistributedSampler这个关键词去pytorch官网中Docs文档中去搜索一下:
在这里插入图片描述
然后点击torch.utils.data.distributed.DistributedSampler就能看到对应参数的介绍,点击SOURCE也可以查看到源码
在这里插入图片描述
DistributedSampler说明
在这里插入图片描述
假设我们当前数据集一共有11个样本,序号为0-10. 通过查看DistributedSampler的源码给的信息知道:第一步会通过shuffle处理打扰数据,假设有2块GPU进行并行计算,首先会通过11/2并进行向上取整,在乘以GPU的个数就等于12.由于只有11个数据,差一个数据,我们就从头开始补数据,将第一个数据补充到最后面,这样就有12个数据,这样就能够均衡的分配到每个GPU设备中。将数据间隔的分配到不同的设备中,标蓝色的表示分配给第一块GPU,标橙色的表示分配给第二块GPU。这样就将整个数据集的数据均分到每块GPU上了。

然后对train_sampler做进一步处理,打包为一个个batch,如果batch_size=2的话,就将数据2个2个一组打包为batch。如果最后的数据不够一个batch,通过设置drop_last=True就将最后不足一个batch的数据给丢掉了。

# 将样本索引每batch_size个元素组成一个listtrain_batch_sampler = torch.utils.data.BatchSampler(   train_sampler, batch_size, drop_last=True)

在这里插入图片描述
dataloader
dataloader中传入了刚刚定义好的train_batch_sampler,对于验证集的dataloader由于没有使用batch_sampler打包为一个个batch,因此直接传入val_sampler

nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers if rank == 0:     print('Using {} dataloader workers every process'.format(nw)) train_loader = torch.utils.data.DataLoader(train_data_set,                batch_sampler=train_batch_sampler,                pin_memory=True,                num_workers=nw,                collate_fn=train_data_set.collate_fn) val_loader = torch.utils.data.DataLoader(val_data_set,              batch_size=batch_size,              sampler=val_sampler,              pin_memory=True,              num_workers=nw,              collate_fn=val_data_set.collate_fn)

3 ) 模型载入预训练权重

# 实例化模型 model = resnet34(num_classes=num_classes).to(device) # 如果存在预训练权重则载入 if os.path.exists(weights_path):     weights_dict = torch.load(weights_path, map_location=device)     load_weights_dict = {k: v for k, v in weights_dict.items()                          if model.state_dict()[k].numel() == v.numel()}     model.load_state_dict(load_weights_dict, strict=False) else:     checkpoint_path = os.path.join(tempfile.gettempdir(), "initial_weights.pt")     # 如果不存在预训练权重,需要将第一个进程中的权重保存,然后其他进程载入,保持初始化权重一致     if rank == 0:         torch.save(model.state_dict(), checkpoint_path)     dist.barrier()     # 这里注意,一定要指定map_location参数,否则会导致第一块GPU占用更多资源     model.load_state_dict(torch.load(checkpoint_path, map_location=device))

判断是否有权重文件,有权重就载入。如果没有权重的话,由于使用多GPU训练我们必须保证每个设备上它初始的权重是一模一样的。如果初始化权重不一样的话,那么各个GPU上所取得的梯度,就不是针对同一组参数。如果我们没有载入预训练权重,我们要保证每块GPU上的初始变量是一模一样的。这里我们会在主进程保存我们模型初始化的权重,然后在不同设备上载入我们在主进程中保留下来的权重文件。这样我们就能够保证,在开始训练时,各个GPU对应的模型权重都是一模一样的

4 ) 冻结权重,同步BN
和单GPU的脚本是一样的,对除fc以外的权重进行冻结。如果不冻结权重的话,需要判断是否使用具有同步功能的BN,如果args.syncBN为True的话,就使用下面的命令同步模型。

# 使用SyncBatchNorm后训练会更耗时       model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

就会将模型中所有的BN替换为具有同步功能的BN了,然后再指认到我们的设备中。如果我们冻结权重,只有全连接层可训练,但全连接层没有BN,这里使用同步的BN是没有意义的。所以这里只有训练所有权重,并且权重中有BN层的时候,才会去考虑使用具有同步功能的BN层。

5 ) 训练
使用torch.nn.parallel.DistributedDataParallel去包装我们的模型,通过这个方法包装后,模型就可以在各个设备间进行通信了。

然后通过遍历获得需要学习更新的参数,指定SGD优化器,初始学习率args,lr,动量和正则项weight_decay,同样设置cosine学习率变化曲线。

# 转为DDP模型 model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) # optimizer pg = [p for p in model.parameters() if p.requires_grad] optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=0.005) # Scheduler https://arxiv.org/pdf/1812.01187.pdf lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf  # cosine scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

通过train_sampler.set_epoch(epoch)通过BatchSampler组合的一个个Batch的数据是不一样的。这样的话不同设备在每一轮获得的数据就不一样。因为打扰数据的随机种子中传入了self.epoch,从而利用epoch实现每个epoch获得的数据是不一样的.
在这里插入图片描述
接下来介绍train_one_epoch与单GPU的train_one_epoch之间的差异。

loss = reduce_value(loss,average = True)

在单GPU中这一行代码是不起作用的,针对多GPU的平均损失应该是所有设备的损失的均值,跟单GPU中单设备的损失不同。通过reduce_value求得所有设备损失的均值。

def reduce_value(value, average=True):    world_size = get_world_size()    if world_size < 2:  # 单GPU的情况        return value    with torch.no_grad():        dist.all_reduce(value)        if average:            value /= world_size        return value

5 ) 验证 evalute

def evaluate(model, data_loader, device):   model.eval()   # 用于存储预测正确的样本个数   sum_num = torch.zeros(1).to(device)   # 在进程0中打印验证进度   if is_main_process():       data_loader = tqdm(data_loader, file=sys.stdout)   for step, data in enumerate(data_loader):       images, labels = data       pred = model(images.to(device))       pred = torch.max(pred, dim=1)[1]       sum_num += torch.eq(pred, labels.to(device)).sum()   # 等待所有进程计算完毕   if device != torch.device("cpu"):       torch.cuda.synchronize(device)   sum_num = reduce_value(sum_num, average=False)   return sum_num.item()

evalue方法中和单GPU不同的地方,就是在预测正确样本个数存在不同。使用多GPU的时候,每个GPU都会计算它自己对应那一部分训练集预测正确的样本个数,最后是需要对所有GPU预测正确个数进行求和的,这样才能求得在验证集上预测正确的样本总和。所以通过reduce_value这个方法来对它求和,这里average就设置为False,为False的话返回来的就是多个设备之间预测的正确数之和,而不是平均了。

 sum_num = reduce_value(sum_num, average=False)

求的预测正确数之和,然后除以val_sampler.total就得到准确率了。最后在主进程中通过tensorboard保存平均损失mean_loss,精度acc,学习率lr, 同时在主进程中保存每个epoch的权重参数

 sum_num = evaluate(model=model,                          data_loader=val_loader,                          device=device)       acc = sum_num / val_sampler.total_size       if rank == 0:           print("[epoch {}] accuracy: {}".format(epoch, round(acc, 3)))           tags = ["loss", "accuracy", "learning_rate"]           tb_writer.add_scalar(tags[0], mean_loss, epoch)           tb_writer.add_scalar(tags[1], acc, epoch)           tb_writer.add_scalar(tags[2], optimizer.param_groups[0]["lr"], epoch)          torch.save(model.module.state_dict(), "./weights/model-{}.pth".format(epoch))

最后还需要注意一点,当我们从头开始训练我们在我们主进程中生成了一个初始化权重保存为临时文件,训练完成好需要将该临时文件移除。最后会调用cleanup去destory进程组,训练完之后将资源释放。

  # 删除临时缓存文件    if rank == 0:        if os.path.exists(checkpoint_path) is True:            os.remove(checkpoint_path)    cleanup()

3. 训练时间对比

使用python -m torch.distributed.launch --nproc_per_node=8 --use_env train_multi_gpu_using_launch.py启动多GPU训练,对于单机多卡的情况nproc_per_nod对应的就是GPU的数量。use_env将GPU信息写入环境变量中

不同GPU数,推理时间对比
在这里插入图片描述

4 .源码

下载地址:https://download.csdn.net/download/weixin_38346042/86768878?spm=1001.2014.3001.5501

来源地址:https://blog.csdn.net/weixin_38346042/article/details/127198216

阅读原文内容投诉

免责声明:

① 本站未注明“稿件来源”的信息均来自网络整理。其文字、图片和音视频稿件的所属权归原作者所有。本站收集整理出于非商业性的教育和科研之目的,并不意味着本站赞同其观点或证实其内容的真实性。仅作为临时的测试数据,供内部测试之用。本站并未授权任何人以任何方式主动获取本站任何信息。

② 本站未注明“稿件来源”的临时测试数据将在测试完成后最终做删除处理。有问题或投稿请发送至: 邮箱/279061341@qq.com QQ/279061341

软考中级精品资料免费领

  • 历年真题答案解析
  • 备考技巧名师总结
  • 高频考点精准押题
  • 2024年上半年信息系统项目管理师第二批次真题及答案解析(完整版)

    难度     813人已做
    查看
  • 【考后总结】2024年5月26日信息系统项目管理师第2批次考情分析

    难度     354人已做
    查看
  • 【考后总结】2024年5月25日信息系统项目管理师第1批次考情分析

    难度     318人已做
    查看
  • 2024年上半年软考高项第一、二批次真题考点汇总(完整版)

    难度     435人已做
    查看
  • 2024年上半年系统架构设计师考试综合知识真题

    难度     224人已做
    查看

相关文章

发现更多好内容

猜你喜欢

AI推送时光机
位置:首页-资讯-后端开发
咦!没有更多了?去看看其它编程学习网 内容吧
首页课程
资料下载
问答资讯