def forward_sequential(self, X, S, L, mask=None):
""" Compute the transformer layer sequentially, for purposes of debugging
TODO: Rewrite this and self.sample() to use a shared iterator
"""
# Prepare node and edge embeddings
V, E, E_idx = self.features(X, L, mask)
h_V = self.W_v(V)
h_E = self.W_e(E)
# Encoder is unmasked self-attention
mask_attend = gather_nodes(mask.unsqueeze(-1), E_idx).squeeze(-1)
mask_attend = mask.unsqueeze(-1) * mask_attend
for layer in self.encoder_layers:
h_EV = cat_neighbors_nodes(h_V, h_E, E_idx)
h_V = layer(h_V, h_EV, mask_V=mask, mask_attend=mask_attend)
# Decoder alternates masked self-attention
mask_attend = self._autoregressive_mask(E_idx).unsqueeze(-1)
mask_1D = mask.view([mask.size(0), mask.size(1), 1, 1])
mask_bw = mask_1D * mask_attend
mask_fw = mask_1D * (1. - mask_attend)
N_batch, N_nodes = X.size(0), X.size(1)
log_probs = torch.zeros((N_batch, N_nodes, 20))
h_S = torch.zeros_like(h_V)
h_V_stack = [h_V] + [
torch.zeros_like(h_V) for _ in range(len(self.decoder_layers))
]
for t in range(N_nodes):
# Hidden layers
E_idx_t = E_idx[:, t:t + 1, :]
h_E_t = h_E[:, t:t + 1, :, :]
h_ES_t = cat_neighbors_nodes(h_S, h_E_t, E_idx_t)
# Stale relational features for future states
h_ESV_encoder_t = mask_fw[:, t:t + 1, :, :] * cat_neighbors_nodes(
h_V, h_ES_t, E_idx_t)
for l, layer in enumerate(self.decoder_layers):
# Updated relational features for future states
h_ESV_decoder_t = cat_neighbors_nodes(h_V_stack[l], h_ES_t,
E_idx_t)
h_V_t = h_V_stack[l][:, t:t + 1, :]
h_ESV_t = mask_bw[:, t:t +
1, :, :] * h_ESV_decoder_t + h_ESV_encoder_t
h_V_stack[l + 1][:,
t, :] = layer(h_V_t,
h_ESV_t,
mask_V=mask[:,
t:t + 1]).squeeze(1)
# Sampling step
h_V_t = h_V_stack[-1][:, t, :]
logits = self.W_out(h_V_t)
log_probs[:, t, :] = F.log_softmax(logits, dim=-1)
# Update
h_S[:, t, :] = self.W_s(S[:, t])
return log_probs
def sample(self, X, L, mask=None, temperature=1.0):
""" Autoregressive decoding of a model """
# Prepare node and edge embeddings
V, E, E_idx = self.features(X, L, mask)
h_V = self.W_v(V)
h_E = self.W_e(E)
# Encoder is unmasked self-attention
mask_attend = gather_nodes(mask.unsqueeze(-1), E_idx).squeeze(-1)
mask_attend = mask.unsqueeze(-1) * mask_attend
for layer in self.encoder_layers:
h_EV = cat_neighbors_nodes(h_V, h_E, E_idx)
h_V = layer(h_V, h_EV, mask_V=mask, mask_attend=mask_attend)
# Decoder alternates masked self-attention
mask_attend = self._autoregressive_mask(E_idx).unsqueeze(-1)
mask_1D = mask.view([mask.size(0), mask.size(1), 1, 1])
mask_bw = mask_1D * mask_attend
mask_fw = mask_1D * (1. - mask_attend)
N_batch, N_nodes = X.size(0), X.size(1)
h_S = torch.zeros_like(h_V)
S = torch.zeros((N_batch, N_nodes), dtype=torch.int64)
h_V_stack = [h_V] + [
torch.zeros_like(h_V) for _ in range(len(self.decoder_layers))
]
for t in range(N_nodes):
# Hidden layers
E_idx_t = E_idx[:, t:t + 1, :]
h_E_t = h_E[:, t:t + 1, :, :]
h_ES_t = cat_neighbors_nodes(h_S, h_E_t, E_idx_t)
# Stale relational features for future states
h_ESV_encoder_t = mask_fw[:, t:t + 1, :, :] * cat_neighbors_nodes(
h_V, h_ES_t, E_idx_t)
for l, layer in enumerate(self.decoder_layers):
# Updated relational features for future states
h_ESV_decoder_t = cat_neighbors_nodes(h_V_stack[l], h_ES_t,
E_idx_t)
h_V_t = h_V_stack[l][:, t:t + 1, :]
h_ESV_t = mask_bw[:, t:t +
1, :, :] * h_ESV_decoder_t + h_ESV_encoder_t
h_V_stack[l + 1][:,
t, :] = layer(h_V_t,
h_ESV_t,
mask_V=mask[:,
t:t + 1]).squeeze(1)
# Sampling step
h_V_t = h_V_stack[-1][:, t, :]
logits = self.W_out(h_V_t) / temperature
probs = F.softmax(logits, dim=-1)
S_t = torch.multinomial(probs, 1).squeeze(-1)
# Update
h_S[:, t, :] = self.W_s(S_t)
S[:, t] = S_t
return S
def forward(self, X, S, L, mask):
""" Graph-conditioned sequence model """
# Prepare node and edge embeddings
V, E, E_idx = self.features(X, L, mask)
h_V = self.W_v(V)
h_E = self.W_e(E)
# Encoder is unmasked self-attention
mask_attend = gather_nodes(mask.unsqueeze(-1), E_idx).squeeze(-1)
mask_attend = mask.unsqueeze(-1) * mask_attend
for layer in self.encoder_layers:
h_EV = cat_neighbors_nodes(h_V, h_E, E_idx)
h_V = layer(h_V, h_EV, mask_V=mask, mask_attend=mask_attend)
# Concatenate sequence embeddings for autoregressive decoder
h_S = self.W_s(S)
h_ES = cat_neighbors_nodes(h_S, h_E, E_idx)
# Build encoder embeddings
h_ES_encoder = cat_neighbors_nodes(torch.zeros_like(h_S), h_E, E_idx)
h_ESV_encoder = cat_neighbors_nodes(h_V, h_ES_encoder, E_idx)
# Decoder uses masked self-attention
mask_attend = self._autoregressive_mask(E_idx).unsqueeze(-1)
mask_1D = mask.view([mask.size(0), mask.size(1), 1, 1])
mask_bw = mask_1D * mask_attend
if self.forward_attention_decoder:
mask_fw = mask_1D * (1. - mask_attend)
h_ESV_encoder_fw = mask_fw * h_ESV_encoder
else:
h_ESV_encoder_fw = 0
for layer in self.decoder_layers:
# Masked positions attend to encoder information, unmasked see.
h_ESV = cat_neighbors_nodes(h_V, h_ES, E_idx)
h_ESV = mask_bw * h_ESV + h_ESV_encoder_fw
h_V = layer(h_V, h_ESV, mask_V=mask)
logits = self.W_out(h_V)
log_probs = F.log_softmax(logits, dim=-1)
return log_probs
def forward(self, S, L, mask=None):
""" Build a representation of each position in a sequence """
# V, E, E_idx = self.features(X, L, mask)
E_idx = self._build_indices(S)
E = self.positional_embeddings(E_idx)
h_E = self.W_e(E)
h_S = self.W_s(S)
h_V = torch.zeros([S.size(0), S.size(1), self.hidden_dim],
dtype=torch.float32)
# Decoder alternates masked self-attention
mask_attend = self._autoregressive_mask(E_idx)
h_S_neighbors = gather_nodes(h_S, E_idx)
for gsa in self.decoder_layers:
h_V = gsa(h_V,
h_E,
E_idx,
mask_V=mask,
h_E_aux=h_S_neighbors,
mask_attend=mask_attend)
logits = self.W_out(h_V)
log_probs = F.log_softmax(logits, dim=2)
return log_probs
def _orientations_coarse(self, X, E_idx, eps=1e-6):
# Pair features
# Shifted slices of unit vectors
dX = X[:, 1:, :] - X[:, :-1, :]
U = F.normalize(dX, dim=-1)
u_2 = U[:, :-2, :]
u_1 = U[:, 1:-1, :]
u_0 = U[:, 2:, :]
# Backbone normals
n_2 = F.normalize(torch.cross(u_2, u_1), dim=-1)
n_1 = F.normalize(torch.cross(u_1, u_0), dim=-1)
# Bond angle calculation
cosA = -(u_1 * u_0).sum(-1)
cosA = torch.clamp(cosA, -1 + eps, 1 - eps)
A = torch.acos(cosA)
# Angle between normals
cosD = (n_2 * n_1).sum(-1)
cosD = torch.clamp(cosD, -1 + eps, 1 - eps)
D = torch.sign((u_2 * n_1).sum(-1)) * torch.acos(cosD)
# Backbone features
AD_features = torch.stack((torch.cos(A), torch.sin(A) * torch.cos(D),
torch.sin(A) * torch.sin(D)), 2)
AD_features = F.pad(AD_features, [0, 0, 1, 2], 'constant', 0)
# Build relative orientations
o_1 = F.normalize(u_2 - u_1, dim=-1)
O = torch.stack((o_1, n_2, torch.cross(o_1, n_2)), 2)
O = O.view(list(O.shape[:2]) + [9])
O = F.pad(O, [0, 0, 1, 2], 'constant', 0)
# DEBUG: Viz [dense] pairwise orientations
# O = O.view(list(O.shape[:2]) + [3,3])
# dX = X.unsqueeze(2) - X.unsqueeze(1)
# dU = torch.matmul(O.unsqueeze(2), dX.unsqueeze(-1)).squeeze(-1)
# dU = dU / torch.norm(dU, dim=-1, keepdim=True)
# dU = (dU + 1.) / 2.
# plt.imshow(dU.data.numpy()[0])
# plt.show()
# print(dX.size(), O.size(), dU.size())
# exit(0)
O_neighbors = gather_nodes(O, E_idx)
X_neighbors = gather_nodes(X, E_idx)
# Re-view as rotation matrices
O = O.view(list(O.shape[:2]) + [3, 3])
O_neighbors = O_neighbors.view(list(O_neighbors.shape[:3]) + [3, 3])
# Rotate into local reference frames
dX = X_neighbors - X.unsqueeze(-2)
dU = torch.matmul(O.unsqueeze(2), dX.unsqueeze(-1)).squeeze(-1)
dU = F.normalize(dU, dim=-1)
R = torch.matmul(O.unsqueeze(2).transpose(-1, -2), O_neighbors)
Q = self._quaternions(R)
# Orientation features
O_features = torch.cat((dU, Q), dim=-1)
# DEBUG: Viz pairwise orientations
# IMG = Q[:,:,:,:3]
# # IMG = dU
# dU_full = torch.zeros(X.shape[0], X.shape[1], X.shape[1], 3).scatter(
# 2, E_idx.unsqueeze(-1).expand(-1,-1,-1,3), IMG
# )
# print(dU_full)
# dU_full = (dU_full + 1.) / 2.
# plt.imshow(dU_full.data.numpy()[0])
# plt.show()
# exit(0)
# print(Q.sum(), dU.sum(), R.sum())
return AD_features, O_features