训练模型中如何增加梯度计算

Terence888

1. class CustomModel(nn.Module):
   
2.     def __init__(self, cfg, config_path=None, pretrained=False):
   
3.         super().__init__()
   
4.         self.cfg = cfg
   
5.         if config_path is None:
   
6.             self.config = AutoConfig.from_pretrained(cfg.model, output_hidden_states=True)
   
7.             self.config.hidden_dropout = 0.
   
8.             self.config.hidden_dropout_prob = 0.
   
9.             self.config.attention_dropout = 0.
   
10.             self.config.attention_probs_dropout_prob = 0.
    
11.             self.config.add_pooling_layer = False
    
12.         else:
    
13.             self.config = torch.load(config_path)
    
14.         if pretrained:
    
15.             self.model = AutoModel.from_pretrained(cfg.model, config=self.config)
    
16.         else:
    
17.             self.model = AutoModel.from_config(self.config)
    
18.         
    
19.         self.model.resize_token_embeddings(len(CFG.tokenizer))
    
20.         if self.cfg.gradient_checkpointing:
    
21.             self.model.gradient_checkpointing_enable()
    
22.         self.fc = nn.Linear(self.config.hidden_size, self.cfg.num_labels)
    
23.         self._init_weights(self.fc)
    
24.         
    
25.     def _init_weights(self, module):
    
26.         if isinstance(module, nn.Linear):
    
27.             module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
    
28.             if module.bias is not None:
    
29.                 module.bias.data.zero_()
    
30.         elif isinstance(module, nn.Embedding):
    
31.             module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
    
32.             if module.padding_idx is not None:
    
33.                 module.weight.data[module.padding_idx].zero_()
    
34.         elif isinstance(module, nn.LayerNorm):
    
35.             module.bias.data.zero_()
    
36.             module.weight.data.fill_(1.0)
    
37.         
    
38.     def feature(self, input_ids, attention_mask, token_type_ids ):
    
39.         outputs = self.model( input_ids, attention_mask, token_type_ids)
    
40.         last_hidden_states = outputs[0]
    
41.         feature = last_hidden_states[:, 0, :] ## CLS token
    
42.         return feature
    
43. 
    
44.     def forward(self,  input_ids, attention_mask, token_type_ids):
    
45.         feature = self.feature( input_ids, attention_mask, token_type_ids)
    
46.         output = self.fc(feature)
    
47.         return output.squeeze(-1)

复制代码

1. # ====================================================
   
2. # inference
   
3. # ====================================================
   
4. def inference_fn(test_loader, model, device):
   
5.     preds = []
   
6.     model.eval()
   
7.     model.to(device)
   
8.     tk0 = tqdm(test_loader, total=len(test_loader))
   
9.     for inputs in tk0:
   
10.         inputs["input_ids"] = inputs["input_ids"].to(device)
    
11.         inputs["attention_mask"] = inputs["attention_mask"].to(device)
    
12.         inputs['token_type_ids'] = inputs['token_type_ids'].to(device)
    
13.         
    
14.         with torch.no_grad():
    
15.             with torch.cuda.amp.autocast(
    
16.                 enabled=True, dtype=torch.float16, cache_enabled=True
    
17.             ):
    
18.                 y_preds = model(inputs["input_ids"], inputs["attention_mask"],inputs['token_type_ids']  )
    
19.         preds.append(y_preds.to("cpu").numpy().astype(np.float32))
    
20.     predictions = np.concatenate(preds)
    
21.     return predictions

复制代码

上述代码不使用梯度更新仅适用于测试集，如何修改可将其用于训练集

FishC

要将上述代码用于训练集，你需要在模型的训练过程中计算梯度并更新模型参数。可以通过以下步骤实现：

1. **设置模型为训练模式**：使用`model.train()`。
2. **定义损失函数和优化器**：选择合适的损失函数（如`nn.CrossEntropyLoss()`）和优化器（如`torch.optim.AdamW()`）。
3. **计算损失并反向传播**：使用`loss.backward()`计算梯度。
4. **更新模型参数**：使用`optimizer.step()`更新参数。
5. **清零梯度**：在每个训练步骤开始前使用`optimizer.zero_grad()`清除上一步的梯度。

以下是如何修改代码以适用于训练集的示例：

1. # 定义损失函数和优化器
   
2. criterion = nn.CrossEntropyLoss()
   
3. optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
   
4. 
   
5. # 训练函数
   
6. def train_fn(train_loader, model, device, optimizer, criterion):
   
7.     model.train()
   
8.     model.to(device)
   
9.     tk0 = tqdm(train_loader, total=len(train_loader))
   
10.     for inputs, labels in tk0:
    
11.         inputs, labels = inputs.to(device), labels.to(device)
    
12.         
    
13.         # 清零梯度
    
14.         optimizer.zero_grad()
    
15.         
    
16.         # 前向传播
    
17.         outputs = model(inputs, inputs, inputs)
    
18.         
    
19.         # 计算损失
    
20.         loss = criterion(outputs, labels)
    
21.         
    
22.         # 反向传播
    
23.         loss.backward()
    
24.         
    
25.         # 更新参数
    
26.         optimizer.step()
    
27.         
    
28.         tk0.set_postfix(loss=loss.item())

复制代码

在训练过程中，确保你传入了正确的`train_loader`，并且`inputs`和`labels`是从数据集中获取的。这样就可以在训练集上进行梯度计算和参数更新了。

以上回复来自 FishC AI 小助理，如未能正确解答您的问题，请继续追问。

Terence888

FishC 发表于 2024-11-3 13:14
要将上述代码用于训练集，你需要在模型的训练过程中计算梯度并更新模型参数。可以通过以下步骤实现：

1. ...

1. import gc,os
   
2. from tqdm.auto import tqdm # visualizing tool for progress
   
3. 
   
4. # They will be used to pick the best model.pt given to the valid loss
   
5. best_model_epoch, valid_loss_values = [],[]
   
6. valid_loss_min = [1] # arbitrary loss I set here
   
7. def train(model,device,train_dataloader,valid_dataloader,epochs,loss_fn,optimizer,metric):
   
8. 
   
9.     for epoch in range(epochs):
   
10.         gc.collect() # memory cleaning垃圾回收机制，减少占用内存
    
11.         model.train()
    
12. 
    
13.         train_loss = 0
    
14.         train_step = 0
    
15.         pbar = tqdm(train_dataloader, total=len(train_dataloader))#tqdm参数是一个iterable
    
16. 
    
17.         for batch, labels in pbar: # you can also write like "for batch in tqdm(train_dataloader"
    
18.             optimizer.zero_grad() # initialize
    
19.             train_step += 1
    
20.             
    
21. 
    
22.             train_input_ids = batch['input_ids'].to(device)#batch是一个字典
    
23.             train_attention_mask = batch['attention_mask'].to(device)
    
24.             train_token_type_ids = batch['token_type_ids'].to(device)
    
25.             train_labels = labels.squeeze().to(device).long()#label真实值long()转化成一维张量
    
26.             
    
27.             # You can refer to the class "TweetsModel" for understand
    
28.             # what would be logits
    
29.             logits = model(train_input_ids, train_attention_mask,train_token_type_ids).to(device)
    
30.             predictions = torch.argmax(logits, dim=1) # get an index from larger one
    
31.             detached_predictions = predictions.detach().cpu().numpy()
    
32.             
    
33.             loss = loss_fn(logits, train_labels)
    
34.             loss.backward()
    
35.             optimizer.step()
    
36.             model.zero_grad()
    
37. 
    
38.             train_loss += loss.detach().cpu().numpy().item()
    
39. 
    
40.             pbar.set_postfix({'train_loss':train_loss/train_step})#设置进度条显示信息
    
41.         pbar.close()
    
42. 
    
43.         with torch.no_grad():
    
44.             model.eval()
    
45. 
    
46.             valid_loss = 0
    
47.             valid_step = 0
    
48.             total_valid_score = 0
    
49. 
    
50.             y_pred = [] # for getting f1_score that is a metric of the competition
    
51.             y_true = []
    
52. 
    
53.             pbar = tqdm(valid_dataloader)
    
54.             for batch,labels in pbar:
    
55.                 valid_step += 1
    
56. 
    
57.                 valid_input_ids = batch['input_ids'].to(device)
    
58.                 valid_attention_mask = batch['attention_mask'].to(device)
    
59.                 valid_token_type_ids = batch['token_type_ids'].to(device)
    
60.                 valid_labels = labels.squeeze().to(device).long()
    
61. 
    
62.                 logits = model(valid_input_ids, valid_attention_mask).to(device)
    
63.                 predictions = torch.argmax(logits, dim=1)
    
64.                 detached_predictions = predictions.detach().cpu().numpy()
    
65.                
    
66.                 loss = loss_fn(logits, valid_labels)
    
67.                 valid_loss += loss.detach().cpu().numpy().item()
    
68. 
    
69.                 y_pred.extend(predictions.cpu().numpy())
    
70.                 y_true.extend(valid_labels.cpu().numpy())
    
71. 
    
72.             valid_loss /= valid_step
    
73.             f1 = f1_score(y_true,y_pred)
    
74. 
    
75.             print(f'Epoch [{epoch+1}/{epochs}] Score: {f1}')
    
76.             print(f'Epoch [{epoch+1}/{epochs}] Valid_loss: {valid_loss}')
    
77. 
    
78.             if valid_loss < min(valid_loss_min):
    
79.                 print('model improved!')
    
80.             else:
    
81.                 print('model not improved')
    
82.    
    
83.             torch.save(model.state_dict(), f'epoch:{epoch+1}_model.pt')#state_dict 是一个字典对象，包含了模型的所有可学习参数（如权重和偏置）及其当前值
    
84.             print('save checkpoint!')
    
85.             valid_loss_min.append(valid_loss)
    
86.             print(f'valid_loss_min:{min(valid_loss_min)}')
    
87. 
    
88.         best_model_epoch.append(f'/kaggle/working/epoch:{epoch+1}_model.pt')
    
89.         valid_loss_values.append(valid_loss)
    
90.         print('='*100)
    
91. 
    
92.     select_best_model() # refer to below function
    
93.     print('Train/Valid Completed!!')
    
94.     del train_dataloader, valid_dataloader # memory cleaning
    
95.     gc.collect()
    
96. 
    
97. def select_best_model():
    
98.     best_model = best_model_epoch[np.array(valid_loss_values).argmin()]
    
99.     os.rename(best_model, best_model.split('.pt')[0] + '_best.pt')#重命名文件

复制代码

上述代码train_dataloader要求是什么形状的？batch是一个字典
有如下报错

1. ---------------------------------------------------------------------------
   
2. ValueError                                Traceback (most recent call last)
   
3. Cell In[128], line 4
   
4.       1 print('Training Start!')
   
5.       2 print('=' * 100)
   
6. ----> 4 train(model,
   
7.       5     device,
   
8.       6     train_dataloader,
   
9.       7     valid_dataloader,
   
10.       8     CFG.epochs,
    
11.       9     loss_fn,
    
12.      10     optimizer,
    
13.      11     metric)
    
14.      13 del model,train_dataloader, valid_dataloader
    
15.      14 gc.collect()
    
16. 
    
17. Cell In[127], line 17, in train(model, device, train_dataloader, valid_dataloader, epochs, loss_fn, optimizer, metric)
    
18.      14 train_step = 0
    
19.      15 pbar = tqdm(train_dataloader, total=len(train_dataloader))#tqdm参数是一个iterable
    
20. ---> 17 for batch, labels in pbar: # you can also write like "for batch in tqdm(train_dataloader"
    
21.      18     optimizer.zero_grad() # initialize
    
22.      19     train_step += 1
    
23. 
    
24. ValueError: too many values to unpack (expected 2)

复制代码