File size: 3,782 Bytes
07fcdfe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | import math
import torch
import torch.nn as nn
class EmbedLayer(nn.Module):
def __init__(self, num_vars: int, num_features: int, num_categs: int, hidden_dim: int):
"""
Embedding layer to represent categorical inputs in continuous space.
Parameters
----------
num_vars : int
Number of nodes in the graph
num_categs : int
Max. number of categories that each variable can take.
hidden_dim : int
Output dimensionality of the embedding layer.
"""
super().__init__()
self.num_vars = num_vars
self.hidden_dim = hidden_dim
self.num_categs = num_categs
self.num_features = num_features
# Gene-wise 1/0 learnable embeddings
# We have num_vars*num_categs*num_features possible embeddings to model.
self.num_embeds = self.num_vars*self.num_categs*self.num_features
self.embedding = nn.Embedding(num_embeddings=self.num_embeds, embedding_dim=self.hidden_dim)
self.embedding.weight.data.mul_(2./math.sqrt(self.num_vars))
self.bias = nn.Parameter(torch.zeros(self.num_vars, self.hidden_dim*self.num_features))
# Tensor for mapping each input to its corresponding embedding range in self.embedding
pos_trans = (torch.arange(self.num_vars*self.num_features, dtype=torch.long) * self.num_categs)
self.register_buffer("pos_trans", pos_trans, persistent=False)
def forward(self, x, topK=None, binarize_intervention=False, binarize_input=True, threshold_input=None):
return self.embed_tensor(x, topK, binarize_intervention, binarize_input, threshold_input)
def embed_tensor(self, x, topK=None, binarize_intervention=False, binarize_input=True, threshold_input=None):
'''
x --> [nnodes*nbatch, num_features]
'''
# Number of variables
pos_trans = self.pos_trans.view(self.num_vars, -1).repeat(int(x.shape[0]/self.num_vars), 1)
# NEED TO CODE BINARIZATION OF PREDICTED PERTURBATION FLAG HERE USING DIFFERENT TOPK
if binarize_intervention:
# Binarize values on U'
uprime = x.view(-1, self.num_vars).clone()
sorts = torch.argsort(uprime, dim=1, descending=True)
row_indices = torch.LongTensor([i for i in range(len(topK)) for e in range(int(topK[i].item()))])
col_indices = torch.LongTensor([e for i in range(len(topK)) for e in range(int(topK[i].item()))])
uprime[:, :] = 0
uprime[row_indices, sorts[row_indices, col_indices]] = 1
uprime = uprime.view(-1)
x_bin = uprime.view(-1, 1)
elif binarize_input:
difference = x - threshold_input.reshape(1, -1)
difference = torch.where(difference >= 0, difference.to(torch.double), torch.inf).float()
percentile = torch.argsort(difference)[:, 0]
percentile = torch.where(percentile != len(threshold_input)-1, percentile, len(threshold_input)-2)
x_bin = percentile.view(-1, 1)
else:
x_bin = x
indices = x_bin + pos_trans
# local embedding
if binarize_intervention: # Adding 'x' directly so the gradient can flow through them to model_2
x_local = (self.embedding(indices.int()) + x.unsqueeze(-1).repeat(1, 1, self.hidden_dim)).view(indices.shape[0], -1)
else:
x_local = self.embedding(indices.int()).view(indices.shape[0], -1)
bias_local = self.bias.repeat(int(x_local.shape[0]/self.num_vars), 1)
x_local = x_local + bias_local
return x_local, x_bin
|