mirror of
https://github.com/wassname/Castor.git
synced 2026-08-20 12:00:37 +08:00
87 lines
3.7 KiB
Python
87 lines
3.7 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from utils import torch_util
|
|
|
|
|
|
class StackBiLSTMMaxout(nn.Module):
|
|
|
|
def __init__(self, h_size=[512, 1024, 2048], d=300, mlp_d=1600, dropout_r=0.1, max_l=60, num_classes=3):
|
|
super().__init__()
|
|
|
|
self.arch = "SSE"
|
|
self.lstm = nn.LSTM(input_size=d, hidden_size=h_size[0],
|
|
num_layers=1, bidirectional=True)
|
|
|
|
self.lstm_1 = nn.LSTM(input_size=(d + h_size[0] * 2), hidden_size=h_size[1],
|
|
num_layers=1, bidirectional=True)
|
|
|
|
self.lstm_2 = nn.LSTM(input_size=(d + (h_size[0] + h_size[1]) * 2), hidden_size=h_size[2],
|
|
num_layers=1, bidirectional=True)
|
|
|
|
self.max_l = max_l
|
|
self.h_size = h_size
|
|
|
|
self.mlp_1 = nn.Linear(h_size[2] * 2 * 4, mlp_d)
|
|
self.mlp_2 = nn.Linear(mlp_d, mlp_d)
|
|
self.sm = nn.Linear(mlp_d, num_classes)
|
|
|
|
self.classifier = nn.Sequential(*[self.mlp_1, nn.ReLU(), nn.Dropout(dropout_r),
|
|
self.mlp_2, nn.ReLU(), nn.Dropout(dropout_r),
|
|
self.sm])
|
|
|
|
def display(self):
|
|
for param in self.parameters():
|
|
print(param.data.size())
|
|
|
|
def forward(self, sent1, sent2, ext_feats=None, word_to_doc_count=None, raw_sent1=None, raw_sent2=None):
|
|
sent1 = sent1.permute(2, 0, 1) # from [B * D * T] to [T * B * D]
|
|
sent2 = sent2.permute(2, 0, 1)
|
|
sent1_lengths = torch.tensor([len(s.split(" ")) for s in raw_sent1])
|
|
sent2_lengths = torch.tensor([len(s.split(" ")) for s in raw_sent2])
|
|
if self.max_l:
|
|
sent1_lengths = sent1_lengths.clamp(max=self.max_l)
|
|
sent2_lengths = sent2_lengths.clamp(max=self.max_l)
|
|
if sent1.size(0) > self.max_l:
|
|
sent1 = sent1[:self.max_l, :]
|
|
if sent2.size(0) > self.max_l:
|
|
sent2 = sent2[:self.max_l, :]
|
|
#p_sent1 = self.Embd(sent1)
|
|
#p_sent2 = self.Embd(sent2)
|
|
sent1_layer1_out = torch_util.auto_rnn_bilstm(self.lstm, sent1, sent1_lengths)
|
|
sent2_layer1_out = torch_util.auto_rnn_bilstm(self.lstm, sent2, sent2_lengths)
|
|
|
|
# Length truncate
|
|
len1 = sent1_layer1_out.size(0)
|
|
len2 = sent2_layer1_out.size(0)
|
|
p_sent1 = sent1[:len1, :, :] # [T, B, D]
|
|
p_sent2 = sent2[:len2, :, :] # [T, B, D]
|
|
|
|
# Using residual connection
|
|
sent1_layer2_in = torch.cat([p_sent1, sent1_layer1_out], dim=2)
|
|
sent2_layer2_in = torch.cat([p_sent2, sent2_layer1_out], dim=2)
|
|
|
|
sent1_layer2_out = torch_util.auto_rnn_bilstm(self.lstm_1, sent1_layer2_in, sent1_lengths)
|
|
sent2_layer2_out = torch_util.auto_rnn_bilstm(self.lstm_1, sent2_layer2_in, sent2_lengths)
|
|
|
|
sent1_layer3_in = torch.cat([p_sent1, sent1_layer1_out, sent1_layer2_out], dim=2)
|
|
sent2_layer3_in = torch.cat([p_sent2, sent2_layer1_out, sent2_layer2_out], dim=2)
|
|
|
|
sent1_layer3_out = torch_util.auto_rnn_bilstm(self.lstm_2, sent1_layer3_in, sent1_lengths)
|
|
sent2_layer3_out = torch_util.auto_rnn_bilstm(self.lstm_2, sent2_layer3_in, sent2_lengths)
|
|
|
|
sent1_layer3_maxout = torch_util.max_along_time(sent1_layer3_out, sent1_lengths)
|
|
sent2_layer3_maxout = torch_util.max_along_time(sent2_layer3_out, sent2_lengths)
|
|
|
|
# Only use the last layer
|
|
features = torch.cat([sent1_layer3_maxout, sent2_layer3_maxout,
|
|
torch.abs(sent1_layer3_maxout - sent2_layer3_maxout),
|
|
sent1_layer3_maxout * sent2_layer3_maxout],
|
|
dim=1)
|
|
|
|
out = self.classifier(features)
|
|
out = F.log_softmax(out, dim=1)
|
|
return out
|