郑州php网站开发培训,wordpress 标题重复,建设教育网站,十堰seo推广如何使用torchrun启动单机多卡DDP并行训练
这是一个最近项目中需要使用的方式#xff0c;新近的数据集大概在40w的规模#xff0c;而且载入的原始特征都比较大#xff08;5#xff5e;7M#xff09;#xff0c;所以准备尝试DistributedDataParallel#xff1b;
主要目…如何使用torchrun启动单机多卡DDP并行训练
这是一个最近项目中需要使用的方式新近的数据集大概在40w的规模而且载入的原始特征都比较大57M所以准备尝试DistributedDataParallel
主要目的有两个
1是通过并行训练有效降低训练耗时之前也使用过单机单卡以及单机多卡DataParallel的方式2是通过使用并行训练的方式倍增有效batch_size的数量提高单卡对显存的利用率仍然是为了加速
类似的实现代码你可以在许多技术博客上检索到以下是我实际实践的主要代码
单机单卡到单机多卡的代码修改
# 导入库
import math
# import torch
from torch.utils.data import Dataset, DataLoader
import torch.distributed as dist ## DDP
from torch.utils.data.distributed import DistributedSampler ## DDP
from torch.nn.parallel import DistributedDataParallel as DDP ## DDP
from torch.distributed import init_process_group, destroy_process_group ## DDP# 设置可见的GPU
# os.environ[CUDA_VISIBLE_DEVICES] 0,1,2,3
# 忽略冗余数据存在的潜在问题
# os.environ[KMP_DUPLICATE_LIB_OK] Truedef train():# setuplocal_rank int(os.environ[LOCAL_RANK]) ## DDP init_process_group(backendnccl)torch.cuda.set_device(local_rank)# ...# local_rank dist.get_rank()total_rank dist.get_world_size()device torch.device(cuda, local_rank)# random resetseed 1234torch.manual_seed(seed)torch.cuda.manual_seed(seed)torch.cuda.manual_seed_all(seed)random.seed(seed)np.random.seed(seed)# model Model().to(device)if local_rank 0:print(model)# model.load(local_model_path)lr * math.sqrt(total_rank)# optimizer ...# scheduler ...## train_datas Dataset(train)sampler DistributedSampler(train_datas)# 或者自定义# sampler MyDistributedSampler(train_data_path)train_dataloader DataLoader(train_datas,batch_sizebatch_size,num_workersnum_workers,shuffle(sampler is None),samplersampler,pin_memoryTrue,drop_lastTrue,)train_dataloaders [train_dataloader]# warp modelmodel torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)model DDP(model, device_ids[local_rank], output_devicelocal_rank, find_unused_parametersTrue)for epoch in range(max_epoch):e_INDEX train_epoch(epoch, local_rank, total_rank, device, train_dataloaders, model, optimizer, scheduler)dist.destroy_process_group()def train_epoch(epoch, local_rank, total_rank, device, train_dataloaders, model, optimizer, scheduler):model.train()if local_rank 0:torch.distributed.barrier()# 如果使用了自定义的Sampler 记得实现一个对应set_epoch方法功能的新方法并由采样器进行调用if local_rank 0:for i, train_dataloader in enumerate(train_dataloaders):# train_dataloader.sampler.set_epoch(epoch)train_dataloader.sampler.shuffle_data_on_ranks(seedepoch i)e_loss 0e_num 0# ...for train_dataloader in train_dataloaders:for data in train_dataloader:loss train_step(data, local_rank, total_rank, device, train_dataloaders, model, optimizer, scheduler)e_loss loss[0]e_num loss[1]if local_rank 0:torch.distributed.barrier()total torch.tensor(copy.deepcopy([e_loss, e_num]),dtypetorch.float32,devicedevice,)torch.distributed.all_reduce(total, torch.distributed.ReduceOp.SUM, async_opFalse)e_loss, e_num total.tolist()e_loss / e_num# if local_rank 0:if isinstance(model, torch.nn.parallel.DistributedDataParallel):model.module.save(opt.notes) else:model.save(opt.notes) scheduler.step(e_loss) # if need e_loss# scheduler.step()# 测试每个GPU进程输出的模型第一层参数是相同的for param in model.parameters():print( GPU{} Model param layer1.format(local_rank), param)breake_INDEX eval_epcoh(model, val_dataloader, epoch, local_rank)return e_INDEXdef train_step(data, local_rank, total_rank, device, train_dataloaders, model, optimizer, scheduler):data_a,label_a datadata_a data_a.to(device)label_a label_a.to(device)# 这里要注意如果使用的是自定义api 如这里的compute_loss该方法中 包含了模型的推理和损失计算过程那么就需要使用model.module进行调用# 如果直接使用模型推理 那么直接调用 forward函数即可如model(data_a)if isinstance(model, torch.nn.parallel.DistributedDataParallel):loss, _ model.module.compute_loss(data_a, label_a)else:loss, _loss_dict_ model.compute_loss(data_a, label_a) optimizer.zero_grad() loss.backward() optimizer.step() return loss.item(), label_a.shape[0]
对于测试集的验证可以不使用Sampler验证集/测试集在每个GPU进程中都相同仅在某一个GPU中完整测试即可
启动方式
torchrun -m --nnodes1 --nproc_per_node4 mainnohup torchrun -m --nnodes1 --nproc_per_node4 main 127.0.0.1 5301 /dev/null 21 disown
测试模型的第一层输出part可以看到它们都是相同的 GPU0 Model param layer1 Parameter containing:
tensor([[[[ 0.1158, 0.0723, -0.0191],[ 0.0715, -0.1204, -0.0295],[-0.0790, 0.0723, 0.0963],GPU3 Model param layer1 Parameter containing:
tensor([[[[ 0.1158, 0.0723, -0.0191],[ 0.0715, -0.1204, -0.0295],[-0.0790, 0.0723, 0.0963],GPU1 Model param layer1 Parameter containing:
tensor([[[[ 0.1158, 0.0723, -0.0191],[ 0.0715, -0.1204, -0.0295],[-0.0790, 0.0723, 0.0963],GPU2 Model param layer1 Parameter containing:
tensor([[[[ 0.1158, 0.0723, -0.0191],[ 0.0715, -0.1204, -0.0295],[-0.0790, 0.0723, 0.0963], 一点经验
其实主要的代码并不复杂但是这里我要提到的一点经验是如果你的模型参数在每个epoch结束之后梯度同步更新完成的打印结果并不一致并且你的DDP主要代码也与上述代码大同小异那么很有可能是你的模型有问题并不是不正确只不过无法得到DDP的支持
这也是我的项目所面临的问题我使用了一个相对简单的模型可以很好的实践DDP但另外一个复杂许多的模型就无法通过测试每个epoch后模型参数在各个GPU进程间会变的不同具体的后续补充
关于MyDistributedSampler
from torch.utils.data.sampler import Samplerclass MyDistributedSampler(Sampler):def __init__(self, data_path, m, batch_size, local_rank, total_rank, device, distributeFalse):# 这里的读取 与 Dataset中 读取原始数据是一致的我们使用的文件是exceldata_lines pd.read_excel(data_path)self.distribute distributeself.total_rank -1self.local_rank -1self.device deviceif distribute:if not dist.is_available():raise RuntimeError(Requires distributed package to be available)self.total_rank total_rankself.local_rank local_rankself._m_per_class mself._batch_size batch_size# 拿到与单卡一致的所有数据self._labels_to_indices self._get_labels_to_indices(data_lines)# 拿到与单卡一致的索引数据self._global_labels (self._labels_to_indices.reset_index()[a_cls].drop_duplicates().values.tolist())if not self.distribute:# 单卡情况 待操作数据 就是 所有的索引数据self.labels self._global_labelselse:# 多卡时则要按照 当前的local_rank 进行划分self.labels self._split_label_randoms(seed0)[self.local_rank]assert (self._batch_size % self._m_per_class) 0, m_per_class must divide batch_size without any remainder# 采样器的样本数self._sample_length self._get_sample_length()print(Init Sampler with Mper with {} items, and m {}, batch_num {}\n.format(self._sample_length, m, self.num_iters()))returndef __iter__(self):# 创建一个迭代器迭代长度 与样本数量一致 并初始化为0idx_list [0] * self._sample_lengthi 0num_iters self.num_iters()for _ in range(num_iters):random.shuffle(self.labels)# 为使用HardTripletLoss计算损失这里随机 固定个数的分类 以batch_size32 m4为例则这里去的个数为8# 对于我们的数据集实际count的最小关联数是6 因此将m设为4是合理的curr_label_set self.labels[: self._batch_size // self._m_per_class]for label in curr_label_set:if len(self._labels_to_indices.loc[label]) self._m_per_class:items self._labels_to_indices.loc[label].sample(self._m_per_class) else:items self._labels_to_indices.loc[label].sample(self._m_per_class,replaceTrue) t []for item in items.index.values:cur_id np.random.choice(self._labels_to_indices.loc[label, item][ids],1,replaceFalse,)[0]t.append(cur_id)# 确定采样规范 并返回可以看到这里已经是一个符合要求的完整序列# 注意对应的 Dataloader 不可以使用shuffleidx_list[i : i self._m_per_class] t[: self._m_per_class]i self._m_per_classreturn iter(idx_list)# set_epoch的实现: 仅多卡的时候才需要调用def shuffle_data_on_ranks(self, seed, display_detailsFalse):# 每个epoch都要重新划分各个GPU进程所分的的数据它们彼此不重复self.labels self._split_label_randoms(seed)[self.local_rank]self._sample_length self._get_sample_length()if display_details:display sorted(self.labels)[:5]print(Shuffle sampler with {} items, rank {}, total rank {}, batch_num {}, label of head{} is {}.format(self._sample_length,self.local_rank,self.total_rank,self.num_iters(),len(display),display,))def num_iters(self):# 样本数 整除 batch_size的 值即整个训练集被迭代一次 所需要的次数return self._sample_length // self._batch_sizedef _get_sample_length(self):# 这就是实际要加载的样本数由于这里的self.labels 已经根据local_rank完成了划分因此这个sample_len就是每个GPU进程所分得得样本数sample_length sum([self._labels_to_indices.loc[k][count].sum() for k in self.labels])# 由于梯度同步时需要每个GPU进程都有数据因此如下这段同步只要是为了 保证所有GPU进程中的采样数量是一致的与最小值相同如果你的数据很规整 这段同步实际用不到if self.total_rank 0:torch.distributed.barrier()total torch.tensor([sample_length],dtypetorch.float32,deviceself.device,)torch.distributed.all_reduce(total, torch.distributed.ReduceOp.MIN, async_opTrue)sample_length int(total.tolist()[0])sample_length - sample_length % self._batch_sizereturn sample_lengthdef _split_label_randoms(self, seed):split_label []global_label self._global_labels.copy()random.Random(seed).shuffle(global_label)for i in range(self.total_rank):# 可以看到 索引数据会根据 total_rank 间隔取值split_label.append(global_label[i :: self.total_rank])return split_label# staticmethoddef _get_labels_to_indices(self, data_lines):df data_lines.groupby([a_cls, a_item]).agg(count(a_id, count), ids(a_id, unique))return dfdef __len__(self):return self._sample_length
