|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
|
|
class Embed(nn.Module): |
|
def __init__(self, input_dim, output_dim, normalize_input=False, event_level=False, activation='gelu'): |
|
super().__init__() |
|
|
|
self.input_bn = nn.BatchNorm1d(input_dim) if normalize_input else None |
|
self.fc1 = nn.Linear(input_dim, output_dim) |
|
self.fc2 = nn.Linear(output_dim, output_dim) |
|
self.fc3 = nn.Linear(output_dim, output_dim) |
|
self.event_level = event_level |
|
|
|
|
|
def forward(self, x): |
|
if self.input_bn is not None: |
|
|
|
x = self.input_bn(x) |
|
if not self.event_level: |
|
x = x.permute(2, 0, 1).contiguous() |
|
|
|
x = F.relu(self.fc1(x)) |
|
x = F.relu(self.fc2(x)) |
|
x = F.relu(self.fc3(x)) |
|
|
|
return x |
|
|
|
|
|
class AttBlock(nn.Module): |
|
def __init__(self, embed_dims, linear_dims1, linear_dims2, num_heads=8, activation='relu'): |
|
super(AttBlock, self).__init__() |
|
|
|
self.layer_norm1 = nn.LayerNorm(embed_dims) |
|
self.multihead_attention = nn.MultiheadAttention(embed_dims, num_heads) |
|
self.layer_norm2 = nn.LayerNorm(embed_dims) |
|
self.linear1 = nn.Linear(embed_dims, linear_dims1) |
|
self.activation = nn.ReLU() if activation == 'relu' else nn.GELU() |
|
self.layer_norm3 = nn.LayerNorm(linear_dims1) |
|
self.linear2 = nn.Linear(linear_dims1, linear_dims2) |
|
|
|
def forward(self, x, padding_mask=None): |
|
|
|
x = self.layer_norm1(x) |
|
|
|
if padding_mask is not None: |
|
|
|
|
|
padding_mask = padding_mask.bool() |
|
|
|
|
|
x_att, attention = self.multihead_attention(x, x, x, key_padding_mask=padding_mask, need_weights=True, average_attn_weights=True) |
|
|
|
|
|
x = x + x_att |
|
|
|
x = self.layer_norm2(x) |
|
|
|
x_linear1 = self.activation(self.linear1(x)) |
|
|
|
x = x + x_linear1 |
|
|
|
x = self.layer_norm3(x_linear1) |
|
|
|
x_linear2 = self.linear2(x) |
|
|
|
x = x + x_linear2 |
|
return x, attention |
|
|
|
class ClassBlock(nn.Module): |
|
def __init__(self, embed_dims, linear_dims1, linear_dims2, num_heads=8, activation='relu'): |
|
super(ClassBlock, self).__init__() |
|
|
|
self.layer_norm1 = nn.LayerNorm(embed_dims) |
|
self.multihead_attention = nn.MultiheadAttention(embed_dims, num_heads) |
|
self.layer_norm2 = nn.LayerNorm(embed_dims) |
|
self.linear1 = nn.Linear(embed_dims, linear_dims1) |
|
self.activation = nn.ReLU() if activation == 'relu' else nn.GELU() |
|
self.layer_norm3 = nn.LayerNorm(linear_dims1) |
|
self.linear2 = nn.Linear(linear_dims1, linear_dims2) |
|
|
|
def forward(self, x, class_token, padding_mask=None): |
|
|
|
x = torch.cat((class_token, x), dim=0) |
|
|
|
x = self.layer_norm1(x) |
|
|
|
|
|
if padding_mask is not None: |
|
|
|
padding_mask = torch.cat((torch.zeros_like(padding_mask[:, :1]), padding_mask), dim=1) |
|
padding_mask = padding_mask.bool() |
|
|
|
|
|
x_att, attention = self.multihead_attention(class_token, x, x, key_padding_mask=padding_mask, need_weights=True, average_attn_weights=False) |
|
|
|
x = self.layer_norm2(x_att) |
|
x = class_token + x |
|
|
|
x_linear1 = self.activation(self.linear1(x)) |
|
|
|
x_linear1 = self.layer_norm3(x_linear1) |
|
|
|
x_linear2 = self.linear2(x_linear1 ) |
|
|
|
x = x + x_linear2 |
|
return x, attention |
|
|
|
class MLPHead(nn.Module): |
|
def __init__(self, input_dim, hidden_dim1, hidden_dim2, output_dim): |
|
super(MLPHead, self).__init__() |
|
self.fc1 = nn.Linear(input_dim, hidden_dim1) |
|
self.fc2 = nn.Linear(hidden_dim1, hidden_dim2) |
|
self.fc3 = nn.Linear(hidden_dim2, output_dim) |
|
|
|
def forward(self, x): |
|
x = F.relu(self.fc1(x)) |
|
x = F.relu(self.fc2(x)) |
|
x = self.fc3(x) |
|
return x |
|
|
|
|
|
class AnalysisObjectTransformer(nn.Module): |
|
def __init__(self, input_dim_obj, input_dim_event, embed_dims, linear_dims1, linear_dims2, mlp_hidden_1, mlp_hidden_2, num_heads=8): |
|
super(AnalysisObjectTransformer, self).__init__() |
|
|
|
self.embed_dims = embed_dims |
|
|
|
|
|
self.embedding_layer = Embed(input_dim_obj, embed_dims) |
|
self.embedding_layer_event_level = Embed(input_dim_event, embed_dims, event_level=True) |
|
|
|
|
|
self.block1 = AttBlock(embed_dims, linear_dims1, linear_dims1, num_heads) |
|
self.block2 = AttBlock(linear_dims1, linear_dims1, linear_dims1, num_heads) |
|
self.block3 = AttBlock(linear_dims1, linear_dims2, linear_dims2, num_heads) |
|
self.block5 = ClassBlock(linear_dims2, linear_dims1, linear_dims2, num_heads) |
|
self.block6 = ClassBlock(linear_dims2, linear_dims1, linear_dims2, num_heads) |
|
self.block7 = ClassBlock(linear_dims2, linear_dims1, linear_dims2, num_heads) |
|
|
|
|
|
|
|
self.mlp = MLPHead(embed_dims + input_dim_event, mlp_hidden_1, mlp_hidden_2, output_dim=1) |
|
self.sigmoid = nn.Sigmoid() |
|
|
|
def forward(self, x, event_level, mask=None): |
|
|
|
cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims), requires_grad=True) |
|
cls_token = nn.init.trunc_normal_(cls_token, std=.02) |
|
|
|
|
|
x = self.embedding_layer(x) |
|
x = x.permute(1, 0, 2) |
|
|
|
attention_weights = [] |
|
|
|
|
|
x, attention = self.block1(x, padding_mask=mask) |
|
|
|
attention_weights.append(attention) |
|
x, attention = self.block2(x, padding_mask=mask) |
|
attention_weights.append(attention) |
|
x, attention = self.block3(x, padding_mask=mask) |
|
attention_weights.append(attention) |
|
|
|
cls_tokens = cls_token.expand(1, x.size(1), -1) |
|
cls_tokens, attention = self.block5(x, cls_tokens, padding_mask=mask) |
|
cls_tokens, attention = self.block6(x, cls_tokens, padding_mask=mask) |
|
cls_tokens, attention = self.block7(x, cls_tokens, padding_mask=mask) |
|
|
|
x = torch.cat((cls_tokens.squeeze(0), event_level), dim=-1) |
|
x = self.mlp(x) |
|
output_probabilities = self.sigmoid(x) |
|
return output_probabilities, attention_weights |
|
|