353 lines
12 KiB
Python
353 lines
12 KiB
Python
import torch.nn as nn
|
|
import torch
|
|
import functools
|
|
from torch_geometric.utils import dense_to_sparse
|
|
|
|
from . import utils, layers, gnns
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from .transformer import Encoder, SemanticEmbedding
|
|
from models.GDSS.layers import MLP
|
|
from .set_encoder.setenc_models import SetPool
|
|
|
|
""" Transformer Encoder """
|
|
class GraphEncoder(nn.Module):
|
|
def __init__(self, config):
|
|
super(GraphEncoder, self).__init__()
|
|
# Forward Transformers
|
|
self.encoder_f = Encoder(config)
|
|
|
|
def forward(self, x, mask):
|
|
h_f, hs_f, attns_f = self.encoder_f(x, mask)
|
|
h = torch.cat(hs_f, dim=-1)
|
|
return h
|
|
|
|
|
|
@staticmethod
|
|
def get_embeddings(h_x):
|
|
h_x = h_x.cpu()
|
|
return h_x[:, -1]
|
|
|
|
class CLSHead(nn.Module):
|
|
def __init__(self, config, init_weights=None):
|
|
super(CLSHead, self).__init__()
|
|
self.layer_1 = nn.Linear(config.d_model, config.d_model)
|
|
self.dropout = nn.Dropout(p=config.dropout)
|
|
self.layer_2 = nn.Linear(config.d_model, config.n_vocab)
|
|
if init_weights is not None:
|
|
self.layer_2.weight = init_weights
|
|
|
|
def forward(self, x):
|
|
x = self.dropout(torch.tanh(self.layer_1(x)))
|
|
return F.log_softmax(self.layer_2(x), dim=-1)
|
|
|
|
|
|
@utils.register_model(name='CATE')
|
|
class CATE(nn.Module):
|
|
def __init__(self, config):
|
|
super(CATE, self).__init__()
|
|
# Shared Embedding Layer
|
|
self.opEmb = SemanticEmbedding(config.model.graph_encoder)
|
|
self.dropout_op = nn.Dropout(p=config.model.dropout)
|
|
self.d_model = config.model.graph_encoder.d_model
|
|
self.act = act = get_act(config)
|
|
# Time
|
|
self.timeEmb1 = nn.Linear(self.d_model, self.d_model * 4)
|
|
self.timeEmb2 = nn.Linear(self.d_model * 4, self.d_model)
|
|
|
|
# 2 GraphEncoder for X and Y
|
|
self.graph_encoder = GraphEncoder(config.model.graph_encoder)
|
|
|
|
self.fdim = int(config.model.graph_encoder.n_layers * config.model.graph_encoder.d_model)
|
|
self.final = MLP(num_layers=3, input_dim=self.fdim, hidden_dim=2*self.fdim, output_dim=config.data.n_vocab,
|
|
use_bn=False, activate_func=F.elu)
|
|
|
|
self.pos_enc_type = config.model.pos_enc_type
|
|
self.pos_encoder = PositionalEncoding_StageWise(d_model=self.d_model, max_len=config.data.max_node)
|
|
|
|
def forward(self, X, time_cond, maskX):
|
|
|
|
# Shared Embeddings
|
|
emb_x = self.dropout_op(self.opEmb(X))
|
|
|
|
if self.pos_encoder is not None:
|
|
emb_p = self.pos_encoder(emb_x) # [20, 64]
|
|
emb_x = emb_x + emb_p
|
|
# Time embedding
|
|
timesteps = time_cond
|
|
emb_t = get_timestep_embedding(timesteps, self.d_model)# time embedding
|
|
emb_t = self.timeEmb1(emb_t) # [32, 512]
|
|
emb_t = self.timeEmb2(self.act(emb_t)) # [32, 64]
|
|
emb_t = emb_t.unsqueeze(1)
|
|
emb = emb_x + emb_t
|
|
|
|
h_x = self.graph_encoder(emb, maskX)
|
|
h_x = self.final(h_x)
|
|
|
|
"""
|
|
Shape: Batch Size, Length (with Pad), Feature Dim (forward) + Feature Dim (backward)
|
|
*HINT: X1 X2 X3 [PAD] [PAD]
|
|
"""
|
|
return h_x
|
|
|
|
|
|
|
|
@utils.register_model(name='PredictorCATE')
|
|
class PredictorCATE(nn.Module):
|
|
def __init__(self, config):
|
|
super(PredictorCATE, self).__init__()
|
|
# Shared Embedding Layer
|
|
self.opEmb = SemanticEmbedding(config.model.graph_encoder)
|
|
self.dropout_op = nn.Dropout(p=config.model.dropout)
|
|
self.d_model = config.model.graph_encoder.d_model
|
|
self.act = act = get_act(config)
|
|
# Time
|
|
self.timeEmb1 = nn.Linear(self.d_model, self.d_model * 4)
|
|
self.timeEmb2 = nn.Linear(self.d_model * 4, self.d_model)
|
|
|
|
# 2 GraphEncoder for X and Y
|
|
self.graph_encoder = GraphEncoder(config.model.graph_encoder)
|
|
|
|
self.fdim = int(config.model.graph_encoder.n_layers * config.model.graph_encoder.d_model)
|
|
self.final = MLP(num_layers=3, input_dim=self.fdim, hidden_dim=2*self.fdim, output_dim=config.data.n_vocab,
|
|
use_bn=False, activate_func=F.elu)
|
|
|
|
self.rdim = int(config.data.max_node * config.data.n_vocab)
|
|
self.regeress = MLP(num_layers=2, input_dim=self.rdim, hidden_dim=2*self.rdim, output_dim=1,
|
|
use_bn=False, activate_func=F.elu)
|
|
|
|
|
|
|
|
def forward(self, X, time_cond, maskX):
|
|
|
|
# Shared Embeddings
|
|
emb_x = self.dropout_op(self.opEmb(X))
|
|
|
|
# Time embedding
|
|
timesteps = time_cond
|
|
emb_t = get_timestep_embedding(timesteps, self.d_model)# time embedding
|
|
emb_t = self.timeEmb1(emb_t) # [32, 512]
|
|
emb_t = self.timeEmb2(self.act(emb_t)) # [32, 64]
|
|
emb_t = emb_t.unsqueeze(1)
|
|
|
|
emb = emb_x + emb_t
|
|
|
|
# h_x = self.graph_encoder(emb_x, maskX)
|
|
h_x = self.graph_encoder(emb, maskX)
|
|
h_x = self.final(h_x)
|
|
|
|
"""
|
|
Shape: Batch Size, Length (with Pad), Feature Dim (forward) + Feature Dim (backward)
|
|
*HINT: X1 X2 X3 [PAD] [PAD]
|
|
"""
|
|
h_x = h_x.reshape(h_x.size(0), -1)
|
|
h_x = self.regeress(h_x)
|
|
|
|
return h_x
|
|
|
|
|
|
class PositionalEncoding_StageWise(nn.Module):
|
|
|
|
def __init__(self, d_model, max_len):
|
|
|
|
super(PositionalEncoding_StageWise, self).__init__()
|
|
|
|
NUM_STAGE = 5
|
|
max_len = int(max_len / NUM_STAGE)
|
|
self.encoding = torch.zeros(max_len, d_model)
|
|
|
|
pos = torch.arange(0, max_len)
|
|
|
|
|
|
pos = pos.float().unsqueeze(dim=1)
|
|
|
|
|
|
_2i = torch.arange(0, d_model, step=2).float()
|
|
|
|
# (max_len, 1) / (d_model/2 ) -> (max_len, d_model/2)
|
|
self.encoding[:, ::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
|
|
self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model))) # (4, 64)
|
|
self.encoding = torch.cat([self.encoding] * NUM_STAGE, dim=0)
|
|
|
|
def forward(self, x):
|
|
batch_size, seq_len, _ = x.size()
|
|
|
|
return self.encoding[:seq_len, :].to(x.device)
|
|
|
|
|
|
@utils.register_model(name='MetaPredictorCATE')
|
|
class MetaPredictorCATE(nn.Module):
|
|
def __init__(self, config):
|
|
super(MetaPredictorCATE, self).__init__()
|
|
|
|
self.input_type= config.model.input_type
|
|
self.hs = config.model.hs
|
|
|
|
# Shared Embedding Layer
|
|
self.opEmb = SemanticEmbedding(config.model.graph_encoder)
|
|
self.dropout_op = nn.Dropout(p=config.model.dropout)
|
|
self.d_model = config.model.graph_encoder.d_model
|
|
self.act = act = get_act(config)
|
|
# Time
|
|
self.timeEmb1 = nn.Linear(self.d_model, self.d_model * 4)
|
|
self.timeEmb2 = nn.Linear(self.d_model * 4, self.d_model)
|
|
|
|
# 2 GraphEncoder for X and Y
|
|
self.graph_encoder = GraphEncoder(config.model.graph_encoder)
|
|
|
|
self.fdim = int(config.model.graph_encoder.n_layers * config.model.graph_encoder.d_model)
|
|
self.final = MLP(num_layers=3, input_dim=self.fdim, hidden_dim=2*self.fdim, output_dim=config.data.n_vocab,
|
|
use_bn=False, activate_func=F.elu)
|
|
|
|
self.rdim = int(config.data.max_node * config.data.n_vocab)
|
|
self.regeress = MLP(num_layers=2, input_dim=self.rdim, hidden_dim=2*self.rdim, output_dim=2*self.rdim,
|
|
use_bn=False, activate_func=F.elu)
|
|
|
|
# Set
|
|
self.nz = config.model.nz
|
|
self.num_sample = config.model.num_sample
|
|
self.intra_setpool = SetPool(dim_input=512,
|
|
num_outputs=1,
|
|
dim_output=self.nz,
|
|
dim_hidden=self.nz,
|
|
mode='sabPF')
|
|
self.inter_setpool = SetPool(dim_input=self.nz,
|
|
num_outputs=1,
|
|
dim_output=self.nz,
|
|
dim_hidden=self.nz,
|
|
mode='sabPF')
|
|
self.set_fc = nn.Sequential(
|
|
nn.Linear(512, self.nz),
|
|
nn.ReLU())
|
|
|
|
input_dim = 0
|
|
if 'D' in self.input_type:
|
|
input_dim += self.nz
|
|
if 'A' in self.input_type:
|
|
input_dim += 2*self.rdim
|
|
|
|
self.pred_fc = nn.Sequential(
|
|
nn.Linear(input_dim, self.hs),
|
|
nn.Tanh(),
|
|
nn.Linear(self.hs, 1)
|
|
)
|
|
|
|
self.sample_state = False
|
|
self.D_mu = None
|
|
|
|
|
|
def arch_encode(self, X, time_cond, maskX):
|
|
# Shared Embeddings
|
|
emb_x = self.dropout_op(self.opEmb(X))
|
|
|
|
# Time embedding
|
|
timesteps = time_cond
|
|
emb_t = get_timestep_embedding(timesteps, self.d_model)# time embedding
|
|
emb_t = self.timeEmb1(emb_t) # [32, 512]
|
|
emb_t = self.timeEmb2(self.act(emb_t)) # [32, 64]
|
|
emb_t = emb_t.unsqueeze(1)
|
|
emb = emb_x + emb_t
|
|
|
|
h_x = self.graph_encoder(emb, maskX)
|
|
h_x = self.final(h_x)
|
|
|
|
h_x = h_x.reshape(h_x.size(0), -1)
|
|
h_x = self.regeress(h_x)
|
|
return h_x
|
|
|
|
def set_encode(self, task):
|
|
proto_batch = []
|
|
for x in task:
|
|
cls_protos = self.intra_setpool(
|
|
x.view(-1, self.num_sample, 512)).squeeze(1)
|
|
proto_batch.append(
|
|
self.inter_setpool(cls_protos.unsqueeze(0)))
|
|
v = torch.stack(proto_batch).squeeze()
|
|
return v
|
|
|
|
def predict(self, D_mu, A_mu):
|
|
input_vec = []
|
|
if 'D' in self.input_type:
|
|
input_vec.append(D_mu)
|
|
if 'A' in self.input_type:
|
|
input_vec.append(A_mu)
|
|
input_vec = torch.cat(input_vec, dim=1)
|
|
return self.pred_fc(input_vec)
|
|
|
|
def forward(self, X, time_cond, maskX, task):
|
|
if self.sample_state:
|
|
if self.D_mu is None:
|
|
self.D_mu = self.set_encode(task)
|
|
D_mu = self.D_mu
|
|
else:
|
|
D_mu = self.set_encode(task)
|
|
A_mu = self.arch_encode(X, time_cond, maskX)
|
|
y_pred = self.predict(D_mu, A_mu)
|
|
return y_pred
|
|
|
|
|
|
class AttentionPool2d(nn.Module):
|
|
"""
|
|
Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
spacial_dim: int,
|
|
embed_dim: int,
|
|
num_heads_channels: int,
|
|
output_dim: int = None,
|
|
):
|
|
super().__init__()
|
|
self.positional_embedding = nn.Parameter(
|
|
torch.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5
|
|
)
|
|
self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
|
|
self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
|
|
self.num_heads = embed_dim // num_heads_channels
|
|
self.attention = QKVAttention(self.num_heads)
|
|
|
|
def forward(self, x):
|
|
b, c, *_spatial = x.shape
|
|
x = x.reshape(b, c, -1) # NC(HW)
|
|
x = torch.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1)
|
|
x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1)
|
|
x = self.qkv_proj(x)
|
|
x = self.attention(x)
|
|
x = self.c_proj(x)
|
|
return x[:, :, 0]
|
|
|
|
import math
|
|
def get_timestep_embedding(timesteps, embedding_dim, max_positions=10000):
|
|
assert len(timesteps.shape) == 1
|
|
half_dim = embedding_dim // 2
|
|
# magic number 10000 is from transformers
|
|
emb = math.log(max_positions) / (half_dim - 1)
|
|
emb = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) * -emb)
|
|
emb = timesteps.float()[:, None] * emb[None, :]
|
|
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
|
|
if embedding_dim % 2 == 1: # zero pad
|
|
emb = F.pad(emb, (0, 1), mode='constant')
|
|
assert emb.shape == (timesteps.shape[0], embedding_dim)
|
|
return emb
|
|
|
|
|
|
def get_act(config):
|
|
"""Get actiuvation functions from the config file."""
|
|
|
|
if config.model.nonlinearity.lower() == 'elu':
|
|
return nn.ELU()
|
|
elif config.model.nonlinearity.lower() == 'relu':
|
|
return nn.ReLU()
|
|
elif config.model.nonlinearity.lower() == 'lrelu':
|
|
return nn.LeakyReLU(negative_slope=0.2)
|
|
elif config.model.nonlinearity.lower() == 'swish':
|
|
return nn.SiLU()
|
|
elif config.model.nonlinearity.lower() == 'tanh':
|
|
return nn.Tanh()
|
|
else:
|
|
raise NotImplementedError('activation function does not exist!')
|