小黑fastNLP成长日记3:model与callback

1.使用torch定义模型

# 使用torch版本
import torch
import torch.nn as nn
class LSTMText(nn.Module):
    def __init__(self,vocab_size,embedding_dim,output_dim,hidden_dim = 64,num_layers = 2,dropout = 0.5):
        super(LSTMText,self).__init__()
        self.embedding = nn.Embedding(vocab_size,embedding_dim)
        self.lstm = nn.LSTM(embedding_dim,hidden_dim,num_layers = num_layers,bidirectional = True,dropout = dropout,batch_first = True)
        self.fc = nn.Linear(hidden_dim * 2,output_dim)
        self.dropout = nn.Dropout(dropout)
    def forward(self,words):
        embedded = self.dropout(self.embedding(words))
        # output:[batch_size,seq_len,hidden_dim*2]
        # hidden:[batch_size,num_layers * 2,hidden_dim]
        # cell:[batch_size,num_layers * 2,hidden_dim]
        output,(hidden,cell) = self.lstm(embedded)
        hidden = torch.cat((hidden[:,-2,:],hidden[:,-1,:]),dim = -1)    # [batch_size,1,2*hidden_size]
        pred = self.fc(hidden.squeeze(1))
        return {'pred':pred}
model = LSTMText(vocab_size = 100,embedding_dim = 287,output_dim = 2)
words = torch.ones([9,10]).long()
model(words)

2.使用fastNLP模块定义模型

from fastNLP.modules import LSTM,MLP
from torch.nn import Embedding
class MyText(nn.Module):
    def __init__(self,vocab_size,embedding_dim,output_dim,hidden_dim = 64,num_layers = 2,dropout = 0.5):
        super(MyText,self).__init__()
        self.embedding = Embedding(vocab_size,embedding_dim)
        self.lstm = LSTM(embedding_dim,hidden_dim,num_layers = num_layers,bidirectional = True)
        self.mlp = MLP([hidden_dim * 2,output_dim],dropout = dropout)
    def forward(self,words):
        embedded = self.embedding(words)
        _,(hidden,_) = self.lstm(embedded)
        pred = self.mlp(torch.cat([hidden[:,-1,:],hidden[:,-2,:]],dim = -1))
        return {'pred':pred}
words = torch.ones([9,10]).long()
model = MyText(vocab_size = 100,embedding_dim = 287,output_dim = 2)
model(words)

3.callback使用

from fastNLP import (Callback,EarlyStopCallback,Trainer,CrossEntropyLoss,AccuracyMetric)
from fastNLP.models import CNNText
import torch.cuda
from fastNLP.io import ChnSentiCorpPipe as pipe
def get_data():
    data = pipe().process_from_file()
    print(data)
    data.rename_field('chars','words')
    train_data = data.get_dataset('train')
    dev_data = data.get_dataset('dev')
    test_data = data.get_dataset('test')
    vocab = data.get_vocab('words')
    tgt_vocab = data.get_vocab('target')
    return train_data,dev_data,test_data,vocab,tgt_vocab
# prepare model
train_data,dev_data,_,vocab,tgt_vocab = get_data()
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
model = CNNText((len(vocab),50),num_classes = len(tgt_vocab))
callbacks = [EarlyStopCallback(5)]
def train_with_callback(cb_list):
    trainer = Trainer(
        device = device,
        n_epochs = 3,
        model = model,
        train_data = train_data,
        dev_data = dev_data,
        loss = CrossEntropyLoss(),
        metrics = AccuracyMetric(),
        callbacks = cb_list,
        check_code_level = -1
    )
    trainer.train()
train_with_callback(callbacks)

4.自定义callback

# 自定义callback
from fastNLP import Callback
from fastNLP import logger

class MyCallBack(Callback):
    def __init__(self):
        super(MyCallBack,self).__init__()
        self.total_loss = 0
        self.start_step = 0
    def on_backward_begin(self,loss):
        self.total_loss += loss.item()
    def on_epoch_end(self):
        n_steps = self.step - self.start_step
        avg_loss = self.total_loss / n_steps
        logger.info('Avg loss at epoch %d,%.6f',self.epoch,avg_loss)
        self.start_step = self.step
callbacks = [MyCallBack()]
train_with_callback(callbacks)

training epochs started 2022-01-02-09-37-11-939445
Evaluate data in 0.97 seconds!
Evaluation on dev at Epoch 1/3. Step:300/900:
AccuracyMetric: acc=0.840833

Avg loss at epoch 1,0.509971
Evaluate data in 0.95 seconds!
Evaluation on dev at Epoch 2/3. Step:600/900:
AccuracyMetric: acc=0.886667

Avg loss at epoch 2,0.849697
Evaluate data in 1.04 seconds!
Evaluation on dev at Epoch 3/3. Step:900/900:
AccuracyMetric: acc=0.901667

Avg loss at epoch 3,1.100345
Reloaded the best model.

In Epoch:3/Step:900, got best dev performance:
AccuracyMetric: acc=0.901667