|
|
from transformers import BertConfig,BertModel
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import torch.nn.functional as F
|
|
|
from torch.autograd import Function
|
|
|
from huggingface_hub import PyTorchModelHubMixin
|
|
|
|
|
|
|
|
|
class CSIBERT(nn.Module):
|
|
|
def __init__(self,bertconfig,input_dim,carrier_attention=False, time_emb=True):
|
|
|
super().__init__()
|
|
|
self.bertconfig=bertconfig
|
|
|
self.auto_pos=time_emb
|
|
|
self.bert=BertModel(bertconfig)
|
|
|
self.hidden_dim=bertconfig.hidden_size
|
|
|
self.input_dim=input_dim
|
|
|
self.carrier_attention=carrier_attention
|
|
|
if carrier_attention:
|
|
|
self.attention = SelfAttention(bertconfig.max_position_embeddings, 128, input_dim)
|
|
|
self.emb=nn.Sequential(
|
|
|
nn.Linear(input_dim, 64),
|
|
|
nn.ReLU(),
|
|
|
nn.Linear(64, self.hidden_dim)
|
|
|
)
|
|
|
|
|
|
else:
|
|
|
self.emb=nn.Sequential(
|
|
|
nn.Linear(input_dim, 64),
|
|
|
nn.ReLU(),
|
|
|
nn.Linear(64, self.hidden_dim)
|
|
|
)
|
|
|
|
|
|
|
|
|
def forward(self,x,attn_mask=None,timestamp=None):
|
|
|
if self.carrier_attention:
|
|
|
x=x.permute(0,2,1)
|
|
|
attn_mat = self.attention(x)
|
|
|
x = torch.bmm(attn_mat, x)
|
|
|
x = x.permute(0, 2, 1)
|
|
|
x=self.emb(x)
|
|
|
if timestamp is not None:
|
|
|
pos_emb=self.positional_embedding(timestamp)
|
|
|
x=x+pos_emb
|
|
|
y=self.bert(inputs_embeds=x,attention_mask=attn_mask, output_hidden_states=True)
|
|
|
y=y.hidden_states[-1]
|
|
|
return y
|
|
|
|
|
|
def mask(self,batch_size=1,min=None,max=None,std=None,avg=None):
|
|
|
if std is not None and avg is not None:
|
|
|
device=std.device
|
|
|
result=torch.randn((batch_size, self.bertconfig.max_position_embeddings ,self.input_dim)).to(device)
|
|
|
result=result*std+avg
|
|
|
else:
|
|
|
result=torch.rand((batch_size, self.bertconfig.max_position_embeddings ,self.input_dim))
|
|
|
if min is not None and max is not None:
|
|
|
device = max.device
|
|
|
result=result.to(device)
|
|
|
result=result*(max-min)+min
|
|
|
return result
|
|
|
|
|
|
def positional_embedding(self,timestamp,t=1):
|
|
|
timestamp**=t
|
|
|
device=timestamp.device
|
|
|
min=torch.min(timestamp,dim=-1,keepdim=True)[0]
|
|
|
max=torch.max(timestamp,dim=-1,keepdim=True)[0]
|
|
|
ran=timestamp.shape[-1]
|
|
|
timestamp=(timestamp-min)/(max-min)*ran
|
|
|
d_model=self.hidden_dim
|
|
|
dim=torch.tensor(list(range(d_model))).to(device)
|
|
|
batch_size,length=timestamp.shape
|
|
|
timestamp=timestamp.unsqueeze(2).repeat(1, 1, d_model)
|
|
|
dim=dim.reshape([1,1,-1]).repeat(batch_size,length,1)
|
|
|
sin_emb = torch.sin(timestamp/10000**(dim//2*2/d_model))
|
|
|
cos_emb = torch.cos(timestamp/10000**(dim//2*2/d_model))
|
|
|
mask=torch.zeros(d_model).to(device)
|
|
|
mask[::2]=1
|
|
|
emb=sin_emb*mask+cos_emb*(1-mask)
|
|
|
return emb
|
|
|
|
|
|
class Token_Classifier(nn.Module):
|
|
|
def __init__(self,bert,class_num=52):
|
|
|
super().__init__()
|
|
|
self.bert=bert
|
|
|
self.classifier=nn.Sequential(
|
|
|
nn.Linear(bert.hidden_dim, 64),
|
|
|
nn.ReLU(),
|
|
|
nn.Linear(64, class_num)
|
|
|
)
|
|
|
|
|
|
def forward(self,x,attn_mask=None,timestamp=None):
|
|
|
x=self.bert(x,attn_mask,timestamp)
|
|
|
x=self.classifier(x)
|
|
|
return x
|
|
|
|
|
|
|
|
|
class GRL(Function):
|
|
|
@staticmethod
|
|
|
def forward(ctx, x, alpha=1):
|
|
|
ctx.alpha = alpha
|
|
|
return x.view_as(x)
|
|
|
|
|
|
@staticmethod
|
|
|
def backward(ctx, grad_output):
|
|
|
output = grad_output.neg() * ctx.alpha
|
|
|
return output, None
|
|
|
|
|
|
|
|
|
class Sequence_Classifier(nn.Module):
|
|
|
def __init__(self,bert,class_num=6):
|
|
|
super().__init__()
|
|
|
self.bert=bert
|
|
|
self.query = nn.Linear(bert.hidden_dim, 64)
|
|
|
self.key = nn.Linear(bert.hidden_dim, 64)
|
|
|
self.value = nn.Linear(bert.hidden_dim, 64)
|
|
|
self.self_attention = nn.MultiheadAttention(embed_dim=64, num_heads=4, dropout=0, batch_first=True)
|
|
|
self.norm1=nn.BatchNorm1d(64)
|
|
|
self.Linear=nn.Linear(64, 64)
|
|
|
self.norm2 = nn.BatchNorm1d(bert.bertconfig.max_position_embeddings * 64)
|
|
|
self.classifier=nn.Sequential(
|
|
|
nn.Linear(bert.bertconfig.max_position_embeddings * 64, 64),
|
|
|
nn.ReLU(),
|
|
|
nn.Linear(64, class_num)
|
|
|
)
|
|
|
self.GRL = GRL()
|
|
|
|
|
|
|
|
|
def forward(self,x,attn_mask=None,timestamp=None,adversarial=False,alpha=1):
|
|
|
x=self.bert(x,attn_mask,timestamp)
|
|
|
if adversarial:
|
|
|
x = self.GRL.apply(x,alpha)
|
|
|
batch_size,length,hidden_dim=x.shape
|
|
|
x_attn, _ = self.self_attention(self.query(x), self.key(x), self.value(x))
|
|
|
x = x + x_attn
|
|
|
x1 = x.reshape(-1, 64)
|
|
|
x1 = self.norm1(x1)
|
|
|
x1 = self.Linear(x1)
|
|
|
x2 = x1.reshape(batch_size, -1)
|
|
|
x2 = self.norm2(x2)
|
|
|
x2=self.classifier(x2)
|
|
|
return x2
|
|
|
|
|
|
class SelfAttention(nn.Module):
|
|
|
def __init__(self, input_dim, da, r):
|
|
|
super().__init__()
|
|
|
self.ws1 = nn.Linear(input_dim, da, bias=False)
|
|
|
self.ws2 = nn.Linear(da, r, bias=False)
|
|
|
|
|
|
def forward(self, h):
|
|
|
attn_mat = F.softmax(self.ws2(torch.tanh(self.ws1(h))), dim=1)
|
|
|
attn_mat = attn_mat.permute(0, 2, 1)
|
|
|
return attn_mat
|
|
|
|
|
|
class Classification(nn.Module):
|
|
|
def __init__(self, csibert, class_num, hs=64, da=128, r=4):
|
|
|
super().__init__()
|
|
|
self.bert = csibert
|
|
|
self.attention = SelfAttention(hs, da, r)
|
|
|
self.classifier = nn.Sequential(
|
|
|
nn.Linear(hs * r, 256),
|
|
|
nn.ReLU(),
|
|
|
nn.Linear(256, class_num)
|
|
|
)
|
|
|
self.GRL = GRL()
|
|
|
|
|
|
|
|
|
def forward(self, x, attn=None, timestamp=None,adversarial=False):
|
|
|
x = self.bert(x, attn, timestamp)
|
|
|
if adversarial:
|
|
|
x = self.GRL.apply(x)
|
|
|
attn_mat = self.attention(x)
|
|
|
m = torch.bmm(attn_mat, x)
|
|
|
flatten = m.view(m.size()[0], -1)
|
|
|
res = self.classifier(flatten)
|
|
|
return res
|
|
|
|
|
|
class CSI_BERT( nn.Module,
|
|
|
PyTorchModelHubMixin
|
|
|
):
|
|
|
def __init__(self, max_len=100, hs=64, layers=4, heads=4, intermediate_size=128, carrier_dim=52, carrier_attn=False, time_embedding=True):
|
|
|
super().__init__()
|
|
|
self.config = BertConfig(max_position_embeddings=max_len, hidden_size=hs, num_hidden_layers=layers,num_attention_heads=heads, intermediate_size=intermediate_size)
|
|
|
self.model = CSIBERT(self.config,carrier_dim,carrier_attn,time_embedding)
|
|
|
|
|
|
|
|
|
def forward(self, x, attn_mask=None, timestamp=None):
|
|
|
return self.model(x,attn_mask,timestamp) |
