单机单卡训练模式
# 设置GPU参数,是否使用GPU,使用那块GPU
if config.use_gpu and torch.cuda.is_available():
device=torch.device('cuda',config.gpu_id)
else:
device=torch.device('cpu')
# 检查一下GPU是否可以使用
print('GPU是否可用:'+str(torch.cuda.is_available()))
单机多卡训练模式
-
Single Machine Data Parallel（单机多卡模式）这个版本已经淘汰
image.png
from torch.nn.parallel import DataParallel
device_id=[0,1,2,3]
device=torch.device('cuda:{}'.format(device_id[0])) # 设置0号GPU是主GPU
model=model.to(device)
model=DataParallel(model,device_ids=device_id,output_device=device)
首先会把所有数据分发到列表上的GPU进行训练，然后再gather到主GPU计算loss
-
DistributedParallel(简称DDP,多进程多卡训练)
image.png
代码变成流程:
- 1.初始化进程组
torch.distributed.init rocess_group(backend="nccl", world_size=n_gpus,rank=args.local_rank)
# backend:进程方式
# word_size:当前这个GPU上有多少张卡
# rank:指定当前进程是在那个GPU上
- 2.设置CUDA_VISIBLE_DEVICES环境变量
torch.cuda.set_device(args.local_rank)
- 3.对模型进行包裹
model = DistributedDataParallel(model.cuda(args.local_rank), device_ids=[args.local_rank])
- 4.对每张卡的数据进行分配
train_sampler = DistributedSampler(train_dataset)
# 源码位于torch/utils/data/distributed.py
- 5.将数据传递到dataload中,传递进来的数据就不需要suffer了
- 6.将数据复制到cuda上
data=data.cuda(args.local_rank)
- 7.执行命令(在使用ddp这种方式训练的时候,需要使用命令执行)
python -m torch.distributed.launch--nproc_per_node=n_gpu train.py
- 8.保存模型
torch.save在local_rank=O的位置进行保存，同样注意调用model.module.state_dict()
torch.load 注意map_location
注意事项:
- train.py中要有接受local_rank的参数选项，launch会传入这个参数
- 每个进程的batch_size应该是一个GPU所需要的batch_size大小
- 在每个周期开始处,调用train_sampler.set_epoch(epoch)可以使得数据充分打乱
- 有了sampler，就不要在DataLoader中设置shuffle=True了
完整代码
# 系统相关的
import argparse
import os
# 框架相关
import torch
from torch.utils.data import DataLoader
import torch.optim as optim
import torch.nn as nn
import os
# 自定义包
from BruceNRE.config import config
from BruceNRE.utils import make_seed,load_pkl
from BruceNRE.process import process
from BruceNRE.dataset import CustomDataset,collate_fn
from BruceNRE import models
from BruceNRE.trainer import train,validate
# 导入分布式训练依赖
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel
__Models__={
"BruceCNN":models.BruceCNN
}
parser=argparse.ArgumentParser(description="关系抽取")
parser.add_argument("--model_name",type=str,default='BruceCNN',help='model name')
parser.add_argument('--local_rank',type=int,default=1,help='local device id on current node')
if __name__=="__main__":
# ====================关键代码==================================
os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"
# 分布式训练初始化
torch.distributed.init_process_group(backend="nccl")
# 设置当前的设备只用这张卡
torch.cuda.set_device(args.local_rank)
# 单机多卡:代表有几块GPU
args.word_size=int(os.getenv("WORLD_SIZE",'1'))
# 获取当前进程的序号,用于进程间的通信
args.global_rank=dist.get_rank()
#=============================================================
model_name=args.model_name if args.model_name else config.model_name
# 为了保证模型每次训练的时候都一样的,设置了一个初始化种子
make_seed(config.seed)
# 数据预处理
process(config.data_path,config.out_path,file_type='csv')
# 加载数据
vocab_path=os.path.join(config.out_path,'vocab.pkl')
train_data_path=os.path.join(config.out_path,'train.pkl')
test_data_path=os.path.join(config.out_path,'test.pkl')
vocab=load_pkl(vocab_path,'vocab')
vocab_size=len(vocab.word2idx)
#CustomDataset是继承了torch.util.data的Dataset类的一个类,用于数据加载,详情见Dataset
train_dataset = CustomDataset(train_data_path, 'train-data')
test_dataset = CustomDataset(test_data_path, 'test-data')
# 测试CNN模型
model=__Models__[model_name](vocab_size,config)
print(model)
#=====================关键代码=================================
# 定义,并且把模型放到GPU上
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
global device
device = torch.device("cuda", local_rank)
# 拷贝模型,将模型放入DistributedDataParallelAPI
model.to(device)
# 加载多GPU
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank,
find_unused_parameters=True)
# 构建一个train-sample
train_sample = DistributedSampler(train_dataset)
test_sample=DistributedSampler(test_dataset)
# 使用分布式训练,一定要把suffle设置为False,因为DistributedSampler会吧数据打乱
train_dataloader = DataLoader(
dataset=train_dataset,
batch_size=config.batch_size,
shuffle=False,
drop_last=True,
collate_fn=collate_fn,
sampler=train_sample
)
test_dataloader = DataLoader(
dataset=test_dataset,
batch_size=config.batch_size,
shuffle=False,
drop_last=True,
collate_fn=collate_fn,
sampler=test_sample
)
# =============================================
# 构建优化器
optimizer=optim.Adam(model.parameters(),lr=config.learing_rate)
scheduler=optim.lr_scheduler.ReduceLROnPlateau(optimizer,'max',factor=config.decay_rate,patience=config.decay_patience)
# 损失函数:交叉熵
loss_fn=nn.CrossEntropyLoss()
# 评价指标,微平均,宏平均
best_macro_f1,best_macro_epoch=0,1
best_micro_f1,best_micro_epoch=0,1
best_macro_model,best_micro_model='',''
print("***************************开始训练*******************************")
for epoch in range(1,config.epoch+1):
train_sample.set_epoch(epoch) # 让每张卡在每个周期中得到的数据是随机的
train(epoch,device,train_dataloader,model,optimizer,loss_fn,config)
macro_f1,micro_f1=validate(test_dataloader,device,model,config)
model_name=model.module.save(epoch=epoch)
scheduler.step(macro_f1)
if macro_f1>best_macro_f1:
best_macro_f1=macro_f1
best_macro_epoch=epoch
best_macro_model=model_name
if micro_f1>best_micro_f1:
best_micro_f1=micro_f1
best_micro_epoch=epoch
best_micro_model=model_name
print("=========================模型训练完成==================================")
print(f'best macro f1:{best_macro_f1:.4f}',f'in epoch:{best_macro_epoch},save in:{best_macro_model}')
print(f'best micro f1:{best_micro_f1:.4f}',f'in epoch:{best_macro_epoch},save_in:{best_micro_model}')
最后在shell后台使用的下面语句运行(暂时我只发现这种方法运行,其他方法还需要寻找)
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py
其中
- torch.distributed.launch表示以分布式的方式启动训练，
- nproc_per_node指定一共就多少个节点，可以设置成显卡的个数
