模型代码片段

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

class Config(object):

    def __init__(self):
        self.pre_bert_path="nghuyong/ernie-1.0"
        self.train_path = 'data/dataset_train.csv'  # 训练集
        self.dev_path = 'data/dataset_valid.csv'  # 验证集
        self.test_path = 'data/test.csv'  # 测试集
        self.class_path = 'data/class.json'  # 类别名单
        self.save_path ='mymodel/ernie.pth'        # 模型训练结果
        self.num_classes=10
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')   # 设备

        self.epochs = 10  # epoch数
        self.batch_size = 128  # mini-batch大小
        self.maxlen = 32  # 每句话处理成的长度(短填长切)
        self.learning_rate = 5e-4                                       # 学习率
        self.hidden_size=768
        self.tokenizer = AutoTokenizer.from_pretrained(self.pre_bert_path)

class Model(nn.Module):
    def __init__(self, config):
        super(Model, self).__init__()
        self.ernie=AutoModel.from_pretrained(config.pre_bert_path)
        #设置不更新预训练模型的参数
        for param in self.ernie.parameters():
            param.requires_grad = False
        self.fc = nn.Linear(config.hidden_size, config.num_classes)
    def forward(self, input):
        out=self.ernie(input_ids =input['input_ids'],attention_mask=input['attention_mask'],token_type_ids=input['token_type_ids'])
        #只取最后一层CLS对应的输出
        out = self.fc(out.pooler_output)
        return out

import json
from mymodel import myBert,myAlbertl,myERNIE
import mydataset
import torch
import pandas as pd
from torch import nn,optim
from torch.utils.data import DataLoader

config=myERNIE.Config()

label_dict=json.load(open(config.class_path,'r',encoding='utf-8'))
# 加载训练，验证，测试数据集
train_df = pd.read_csv(config.train_path)
#这里将标签转化为数字
train_ds=mydataset.GetLoader(train_df['review'],[label_dict[i] for i in train_df['cat']])
train_dl=DataLoader(train_ds,batch_size=config.batch_size,shuffle=True)
valid_df = pd.read_csv(config.dev_path)
valid_ds=mydataset.GetLoader(valid_df['review'],[label_dict[i] for i in valid_df['cat']])
valid_dl=DataLoader(valid_ds,batch_size=config.batch_size,shuffle=True)
test_df = pd.read_csv(config.test_path)
test_ds=mydataset.GetLoader(test_df['review'],[label_dict[i] for i in test_df['cat']])
test_dl=DataLoader(test_ds,batch_size=config.batch_size,shuffle=True)

#计算准确率
def accuracys(pre,label):
    pre=torch.max(pre.data,1)[1]
    accuracy=pre.eq(label.data.view_as(pre)).sum()
    return accuracy,len(label)

#导入网络结构
model=myERNIE.Model(config).to(config.device)

#训练
criterion=nn.CrossEntropyLoss()
optimizer=optim.Adam(model.parameters(),lr=config.learning_rate)
best_loss=float('inf')
for epoch in range(config.epochs):
    train_acc = []
    for batch_idx,(data,target)in enumerate(train_dl):
        inputs = config.tokenizer(list(data),truncation=True, return_tensors="pt",padding=True,max_length=config.maxlen)
        model.train()
        out = model(inputs)
        loss=criterion(out,target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_acc.append(accuracys(out,target))
        train_r = (sum(tup[0] for tup in train_acc), sum(tup[1] for tup in train_acc))
        print('当前epoch:{}\t[{}/{}]{:.0f}%\t损失：{:.6f}\t训练集准确率：{:.2f}%\t'.format(
            epoch, batch_idx, len(train_dl), 100. * batch_idx / len(train_dl), loss.data,
                   100. * train_r[0].numpy() / train_r[1]
        ))
        #每100批次进行一次验证
        if batch_idx%100==0 and batch_idx!=0:
            model.eval()
            val_acc=[]
            loss_total=0
            with torch.no_grad():
                for (data,target) in valid_dl:
                    inputs = config.tokenizer(list(data), truncation=True, return_tensors="pt", padding=True,
                                              max_length=config.maxlen)
                    out = model(inputs)
                    loss_total = criterion(out, target).data+loss_total
                    val_acc.append(accuracys(out,target))
            val_r = (sum(tup[0] for tup in val_acc), sum(tup[1] for tup in val_acc))
            print('损失：{:.6f}\t验证集准确率：{:.2f}%\t'.format(loss_total/len(valid_dl),100. * val_r[0].numpy() / val_r[1]))
            #如果验证损失低于最好损失，则保存模型
            if loss_total < best_loss:
                best_loss = loss_total
                torch.save(model.state_dict(), config.save_path)

#测试
model.load_state_dict(torch.load(config.save_path))
model.eval()
test_acc=[]
with torch.no_grad():
    for (data, target) in test_dl:
        inputs = config.tokenizer(list(data),truncation=True, return_tensors="pt",padding=True,max_length=config.maxlen)
        out = model(inputs)
        test_acc.append(accuracys(out, target))
test_r = (sum(tup[0] for tup in test_acc), sum(tup[1] for tup in test_acc))

print('测试集准确率：{:.2f}%\t'.format(100. * test_r[0].numpy() / test_r[1]))