Pytorch_MyModel_CNN.py

python

# 定义模型结构
from transformers import AutoModel
import torch.nn as nn

class MyModel_CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.plm = AutoModel.from_pretrained('bert-base-uncased')
        self.convs = nn.ModuleList([
            nn.Conv2d(1, 64, (4, 768)),
            nn.Conv2d(1, 64, (8, 768)),
            nn.Conv2d(1, 64, (16, 768)),
        ])
        self.fc = nn.Linear(64*3, 2)
    
    def conv_and_pool(self, x, conv): # x.shape = [8, 1, 200, 768]
        x = torch.relu(conv(x)).squeeze(3) # conv(x).shape = [8, 64, XX, 1]; res.shape = [8, 64, XX]
        x = torch.max_pool1d(x, x.size(2)).squeeze(2) # res.shape = [8, 64]
        return x

    def forward(self, batch_inputs): # batch_inputs.shape = [8, 200]
        encoder_out = self.plm(batch_inputs).last_hidden_state
        out = encoder_out.unsqueeze(1)
        
        out0 = self.conv_and_pool(out, self.convs[0])
        out1 = self.conv_and_pool(out, self.convs[1])
        out2 = self.conv_and_pool(out, self.convs[2])
        
        # 串联操作？这个应该是并联操作才对；
        out = torch.cat([out0, out1, out2], dim=1)       
        
        # out = torch.cat([self.conv_and_pool(out, conv) for conv in self.convs])
        out = self.fc(out)
        return out