问题描述:如题,在用预训练Bert进行分类微调时,用nvidia-smi看显存占用和gpu使用率都接近全满,但是速度和在cpu上训练是一样的。
torch.manual_seed(0)
a=CustomDataset(query_dic=data,table_dic=data2,batch_size=4,tokenizer=tokenizer)
checkpoint = torch.load('model_checkpoint0.pth')
model = classifier(cwd + bert_path, 7)
optimizer = optim.Adam(model.parameters(), lr=5e-6)
model.load_state_dict(checkpoint["net"])
optimizer.load_state_dict(checkpoint["optimizer"])
start_epoch=checkpoint["epoch"]
run(1,a,model,optimizer,start_epoch)
首先实例化模型和优化器,加载checkpoint,将模型、优化器和数据集传入到训练函数run中。
函数run的代码如下:
def run(epochs,data_iter,model,optimizer,start_epoch):
device=torch.device("cuda")
model.to(device)#将模型转移到gpu上
criterion=torch.nn.CrossEntropyLoss(ignore_index=10)
for state in optimizer.state.values(): ##将cpu上的optimizer参数转移到gpu上
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
model.train()
for epoc in range(epochs):
for item in tqdm(data_iter):
x = item[0].to(device)#数据集转移至gpu上
att = item[1].to(device)
sel = item[3].to(device)
new = item[4].to(device)
logit = model(x, att, new)
loss = criterion(logit.permute(0, 2, 1), sel)
optimizer.zero_grad()
loss.backward()
optimizer.step()
tqdm.write(f"epoch : {epoc+start_epoch} | loss : {loss:.5f} ")
del x
del att
del sel
del new
del loss
#torch.cuda.empty_cache()
可以看到数据,模型参数,以及加载的优化器参数,都被转移到了gpu上。
显存占用和gpu使用率都满了。
训练速度是每秒一个iter的样子,和之前在cpu上训练的速度比没有提高。
解决方案:将checkpoint作为参数传入到函数中,模型、优化器的实例化和加载改为在函数中完成。
def run(epochs,data_iter,check_point=None):
device=torch.device("cuda")
model=classifier(cwd+bert_path,7)
optimizer=optim.Adam(model.parameters(),lr=5e-6)
model.to(device)
criterion=torch.nn.CrossEntropyLoss(ignore_index=10)
if check_point is not None:
model.load_state_dict(check_point["net"])
optimizer.load_state_dict(check_point["optimizer"])
for state in optimizer.state.values(): ##将cpu上的optimizer参数转移到cuda上
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
start_epoch=check_point["epoch"]
之后速度就达到了每秒钟9个iter的样子