import torch import torch.nn as nn import torch.optim as optim from tqdm import tqdm from datasets import load_dataset from transformers import BertTokenizer, BertModel from torch.utils.data import DataLoader # 0.测试load_dataset方法 def dm01(): # data = load_dataset('csv', data_files='../data/train.csv') # print("data-->", data) # data = load_dataset('csv', data_files='../data/train.csv', split='train') # print("data-->", data) # 注意：split，只能是train data = load_dataset('csv', data_files='../data/test.csv', split='train') print("data-->", data) print(data[0:4]) # 加载预训练模型 tokenizer = BertTokenizer.from_pretrained("../models/bert-base-chinese") # 这里导入模型要使用BertModel，而不是BertModelForSequenceClassification # 因为，这里使用预训练模型进行特征提取，下游自己指定多分类输出层（线性层） bert_model = BertModel.from_pretrained("../models/bert-base-chinese") # 1.数据预处理 def collate_fn1(batch): # 打印数据 # print("batch-->", batch) # 从样本中提取文本和标签 sentences = [row['text'] for row in batch] labels = [row['label'] for row in batch] # print("sentences-->", sentences) # print("labels-->", labels) # 对文本进行数值化+张量化，对标签进行张量化 data = tokenizer.batch_encode_plus(sentences, padding='max_length', max_length=300, truncation=True, return_tensors='pt') # print("data-->", data) # input_ids：数据数值化+张量化后的结果 # token_type_ids：文本（句子）类型，0代表一个句子 # attention_mask：注意力掩码，0代表填充的，1代表有值（需要掩码） input_ids = data['input_ids'] token_type_ids = data['token_type_ids'] attention_mask = data['attention_mask'] # 对label进行张量化 labels = torch.tensor(labels) # 返回数据 return input_ids, token_type_ids, attention_mask, labels def load_data(): # 使用load_dataset加载数据，得到的就是一个Dataset对象 mydataset = load_dataset('csv', data_files='../data/train.csv', split='train') # 把Dataset转换为DataLoader对象 # collate_fn：对数据进行二次处理的函数，就是这一批次的数据，需要接受一个函数地址 mydataloader = DataLoader(dataset=mydataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=collate_fn1) # print("mydataloader-->", mydataloader) # 遍历 for input_ids, token_type_ids, attention_mask, labels in mydataloader: print("input_ids-->", input_ids) print("token_type_ids-->", input_ids) print("attention_mask-->", input_ids) print("labels-->", input_ids) break # 2.模型构建 class MyModel(nn.Module): def __init__(self): super().__init__() # 自定义下游输出层（线性层） # in_features：输入维度，就是Bert模型的768维度 # out_features：输出类别，也就是下游的任务类别，这里是2分类 self.linear = nn.Linear(in_features=768, out_features=2) def forward(self, input_ids, token_type_ids, attention_mask): # 先试用Bert模型进行特征提取(可以冻结，也可以不冻结，如果效果不好，就不冻结) with torch.no_grad(): out = bert_model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) # print("out1-->", out) # print("out2-->", out.last_hidden_state.shape) # [4, 300, 768] # ：代表取4句话 # 0：300个token中，取第0个，也就是CLS # out.last_hidden_state[:, 0] # out.last_hidden_state[:, 0, :] # print("out3-->", out.pooler_output.shape) # [4, 768] # 然后把特征送给下游的自定义输出线性层，实现二分类任务 output = self.linear(out.pooler_output) # print("output-->", output, output.shape) # [4, 2] return output def test_model(): # 数据 mydataset = load_dataset('csv', data_files='../data/train.csv', split='train') mydataloader = DataLoader(dataset=mydataset, batch_size=4, shuffle=True, drop_last=True, collate_fn=collate_fn1) # 模型 model = MyModel() # 遍历 for input_ids, token_type_ids, attention_mask, labels in mydataloader: model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) break # 3.模型训练 def train_model(): # 3+2+5 # 3：模型，损失函数、优化器 # 2：内外两层循环 # 5：内层循环中，5个步骤：前向传播、损失计算、梯度清零、反向传播、梯度更新 # 模型 model = MyModel() # 损失函数 criterion = nn.CrossEntropyLoss() # 优化器 optimizer = optim.AdamW(params=model.parameters(), lr=5e-4) # 轮次 epochs = 1 # 冻结预训练模型的参数(Bert模型) # 双重保险 for param in bert_model.parameters(): param.requires_grad_(False) # 设置模型为训练模型 model.train() # 加载数据 mydataset = load_dataset('csv', data_files='../data/train.csv', split='train') # 把数据转换为DataLoader mydataloader = DataLoader(dataset=mydataset, shuffle=True, drop_last=True, batch_size=4, collate_fn=collate_fn1) # 外层变量轮次 for epoch in range(epochs): # 内层变量批次 # tqdm：进度条，pip install tqdm for i, (input_ids, token_type_ids, attention_mask, labels) in enumerate(tqdm(mydataloader), start=1): # 前向传播 output = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) # output.shape->[4,2] # 损失计算 loss = criterion(output, labels) # print("loss-->", loss) # 梯度清零 optimizer.zero_grad() # 反向传播 loss.backward() # 梯度更新 optimizer.step() # 打印日志 if i % 5 == 0: pred = torch.argmax(output, dim=-1) accuracy = (pred == labels).sum().item() / len(labels) print("epoch:", epoch + 1, "loss:", loss.item(), "accuracy:", accuracy) # 模型保存 torch.save(model.state_dict(), '../save_model/model_%d.pth' % (epoch + 1)) # 4.模型评估 def model_eval(): # 加载数据 mydataset = load_dataset("csv", data_files='../data/test.csv', split='train') # 转换为DataLoader mydataloader = DataLoader(dataset=mydataset, shuffle=False, drop_last=True, batch_size=4, collate_fn=collate_fn1) # 准备模型 model = MyModel() # 加载刚刚预训练好的模型参数 model.load_state_dict(torch.load('../save_model/model_1.pth')) # 准备参数：total，accuracy num = 0 total = 0 # 设置模型为评估模式 model.eval() # 遍历 for i, (input_ids, token_type_ids, attention_mask, labels) in enumerate(tqdm(mydataloader), start=1): # 把数据送给模型，得到预测值 output = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask) # 基于预测结果，得到预测值 y_pred = torch.argmax(output, dim=-1) num += (y_pred == labels).sum().item() total += len(labels) # 计算准确率 print("accuracy-->", num / total) if __name__ == '__main__': # dm01() # load_data() # test_model() # train_model() model_eval() pass