def full_adj_nll(ei, z):
A = to_dense_adj(ei, max_num_nodes=z.size(0))[0]
A_tilde = z @ z.T
temp_size = A.size(0)
temp_sum = A.sum()
posw = float(temp_size * temp_size - temp_sum) / temp_sum
norm = temp_size * temp_size / float(
(temp_size * temp_size - temp_sum) * 2)
nll_loss_mat = F.binary_cross_entropy_with_logits(input=A_tilde,
target=A,
pos_weight=posw,
reduction='none')
nll_loss = -1 * norm * torch.mean(nll_loss_mat, dim=[0, 1])
return -nll_loss
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
z, _ = to_dense_batch(x, batch)
adj = to_dense_adj(edge_index, batch)
batch_size, num_nodes, _ = z.size()
'''
>>> As = torch.randn(3,2,5)
>>> Bs = torch.randn(3,5,4)
>>> torch.einsum('bij,bjk->bik', As, Bs) # batch matrix multiplication
s = F.sigmoid(torch.einsum('bij,bji->bii', x, torch.transpose(x))) # batch matrix multiplication
'''
return None, x, torch.sigmoid(torch.matmul(z, z.transpose(1, 2))), adj
def eval(self, epoch):
total_closs = 0.0
total_laploss = 0.0
total_nloss = 0.0
total_eloss = 0.0
total_loss = 0.0
num_iters = int(
math.ceil(self.params.val_data_size / self.params.batch_size))
loss = 0.0
self.model.eval()
for i in range(num_iters):
self.optimizer.zero_grad()
gt_vertices, gt_normals, gt_edges, gt_image_feats, proj_gt = next(
self.val_generator)
gt_vertices = torch.Tensor(gt_vertices).type(
dtypeF).requires_grad_(False)
gt_normals = torch.Tensor(gt_normals).type(dtypeF).requires_grad_(
False)
gt_image_feats = torch.Tensor(gt_image_feats).type(
dtypeF).requires_grad_(False)
x, c = self.model.forward(gt_image_feats, gt_vertices, gt_normals)
total_closs += self.model.closs / num_iters
total_laploss += self.model.laploss / num_iters
total_nloss += self.model.nloss / num_iters
total_eloss += self.model.eloss / num_iters
total_loss += self.model.loss / num_iters
print(
f'Validation Epoch: {epoch}, Val Epoch Loss: {total_loss}, Val Epoch CLoss: {total_closs}, Val Epoch NLoss: {total_nloss}, Val Epoch ELoss: {total_eloss}, Val EpochLapLoss: {total_laploss}'
)
# proj_pred = utils.flatten_pred_batch(utils.scaleBack(c.x), A, self.params)
utils.drawPolygons(utils.scaleBack(c.x),
utils.scaleBack(gt_vertices[0]),
gt_edges[0],
proj_pred=None,
proj_gt=None,
color='red',
out=self.params.expt_res_dir + '/../val_out.png',
A=to_dense_adj(c.edge_index).cpu().numpy()[0])
self.model.train()
def __call__(self, data):
adj_mat = np.array(to_dense_adj(data.edge_index)[0])
graph = networkx.convert_matrix.from_numpy_matrix(adj_mat)
if not nx.is_connected(graph):
graph = to_connected(graph)
self.model.fit([graph])
hks = self.model.get_embedding().reshape(eval_points, -1).transpose()
#hks=torch.index_select(torch.from_numpy(hks), 1, indice)
num = self.num_node - len(hks)
m = torch.nn.ConstantPad2d((0, 0, 0, num), 0)
mask = torch.cat((torch.ones(
(len(hks), eval_points)), torch.zeros((num, eval_points))), 0)
hks = m(torch.from_numpy(hks)).float()[:, None, :]
data.emb = hks
data.mask = mask
return data
def forward(self, h, g):
if self.ptb == False:
g = to_dense_adj(g)
g = torch.squeeze(g)
# h = self.s_gcn(g, h)
h = torch.squeeze(h)
h = self.g_unet_forward(g, h)
# classify
# h = self.out_drop(h)
# h = self.out_l_1(h)
# h = self.c_act(h)
# h = self.out_drop(h)
# h = self.out_gcn(g, h)
# h = self.gunet(h, g)
return h
def forward(self, x, edge_index, edge_weight):
adj_mat = to_dense_adj(edge_index, edge_attr=edge_weight)
adj_mat = adj_mat.reshape(adj_mat.size(1), adj_mat.size(2))
deg_out = torch.matmul(adj_mat, torch.ones(size=(adj_mat.size(0), 1)))
deg_out = deg_out.flatten()
deg_in = torch.matmul(torch.ones(size=(1, adj_mat.size(0))), adj_mat)
deg_in = deg_in.flatten()
deg_out_inv = torch.reciprocal(deg_out)
deg_in_inv = torch.reciprocal(deg_in)
row, col = edge_index
norm_out = deg_out_inv[row]
norm_in = deg_in_inv[row] # row for W^T
Tx_0 = x
Tx_1 = x
out = torch.matmul(Tx_0, (self.weight[0])[0]) + torch.matmul(
Tx_0, (self.weight[1])[0])
# propagate_type
if self.weight.size(1) > 1:
Tx_1_o = self.propagate(edge_index, x=x, norm=norm_out, size=None)
Tx_1_i = self.propagate(edge_index, x=x, norm=norm_in, size=None)
out = out + torch.matmul(Tx_1_o,
(self.weight[0])[1]) + torch.matmul(
Tx_1_i, (self.weight[1])[1])
for k in range(2, self.weight.size(1)):
Tx_2_o = self.propagate(edge_index,
x=Tx_1_o,
norm=norm_out,
size=None)
Tx_2_o = 2. * Tx_2_o - Tx_0
Tx_2_i = self.propagate(edge_index,
x=Tx_1_i,
norm=norm_in,
size=None)
Tx_2_i = 2. * Tx_2_i - Tx_0
out = out + torch.matmul(Tx_2_o,
(self.weight[0])[k]) + torch.matmul(
Tx_2_i, (self.weight[1])[k])
Tx_0, Tx_1_o, Tx_1_i = Tx_1, Tx_2_o, Tx_2_i
if self.bias is not None:
out += self.bias
return out
def forward(
self,
x: Union[Tensor, PairTensor],
batch: Union[OptTensor, Optional[PairTensor]] = None) -> Tensor:
""""""
is_bipartite: bool = True
if isinstance(x, Tensor):
x: PairTensor = (x, x)
is_bipartite = False
if x[0].dim() != 2:
raise ValueError("Static graphs not supported in 'GravNetConv'")
b: PairOptTensor = (None, None)
if isinstance(batch, Tensor):
b = (batch, batch)
elif isinstance(batch, tuple):
assert batch is not None
b = (batch[0], batch[1])
# embed the inputs before message passing
msg_activations = self.lin_p(x[0])
# transform to the space dimension to build the graph
s_l: Tensor = self.lin_s(x[0])
s_r: Tensor = self.lin_s(x[1]) if is_bipartite else s_l
edge_index = knn(s_l, s_r, self.k, b[0], b[1]).flip([0])
# edge_index = knn_graph(s_l, self.k, b[0], b[1]).flip([0])
edge_weight = (s_l[edge_index[0]] - s_r[edge_index[1]]).pow(2).sum(-1)
edge_weight = torch.exp(-10.0 *
edge_weight) # 10 gives a better spread
# return the adjacency matrix of the graph for lrp purposes
A = to_dense_adj(
edge_index.to("cpu"),
edge_attr=edge_weight.to("cpu"))[0] # adjacency matrix
# message passing
out = self.propagate(edge_index,
x=(msg_activations, None),
edge_weight=edge_weight,
size=(s_l.size(0), s_r.size(0)))
return self.lin_out(out), A, msg_activations
def __call__(self, data):
adj_mat = np.array(to_dense_adj(data.edge_index)[0])
graph = networkx.convert_matrix.from_numpy_matrix(adj_mat)
if not nx.is_connected(graph):
graph=to_connected(graph)
data.x=create_node_feature_matrix(graph)
data.x=torch.tensor(data.x)
#print(data.x)
self.model.fit(graph,data.x)
data.emb = torch.from_numpy(self.model.get_embedding()).float()
num = self.num_node - len(data.emb)
m = nn.ConstantPad2d((0, 0, 0, num), 0)
mask = torch.cat((torch.ones((len(data.emb), eval_points)), torch.zeros((num, eval_points))), 0)
data.emb = m(data.emb).float()
data.mask = mask
data.emb = data.emb.reshape(-1,20, eval_points)
return data
def test(model, loader, loss_func):
total_loss = 0
model.eval()
test_log = TestLog()
for batch_id, rxn_batch in tqdm(enumerate(loader)):
rxn_batch = rxn_batch.to(model.device)
D_pred, mask, W = model(rxn_batch)
D_gt = to_dense_adj(rxn_batch.edge_index, rxn_batch.batch, rxn_batch.y)
batch_loss = loss_func(D_pred, D_gt) / mask.sum()
total_loss += batch_loss.item()
test_log.add_D(D_pred.detach().cpu().numpy())
test_log.add_W(W.detach().cpu().numpy())
RMSE = math.sqrt(total_loss / len(loader.dataset))
return RMSE, test_log
def __init__(self,
dataset='Cora',
lr=0.01,
weight_decay=5e-4,
max_layer=10,
batch_size=128,
policy=""):
device = 'cpu'
dataset = dataset
path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data',
dataset)
dataset = Planetoid(path, dataset, T.NormalizeFeatures())
data = dataset[0]
adj = to_dense_adj(data.edge_index).numpy()[0]
norm = np.array([np.sum(row) for row in adj])
self.adj = (adj / norm).T
self.init_k_hop(max_layer)
self.model, self.data = Net(max_layer,
dataset).to(device), data.to(device)
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr,
weight_decay=weight_decay)
train_mask = self.data.train_mask.to('cpu').numpy()
self.train_indexes = np.where(train_mask == True)[0]
self.batch_size = len(self.train_indexes) - 1
self.i = 0
self.val_acc = 0.0
self._set_action_space(max_layer)
obs = self.reset()
self._set_observation_space(obs)
self.policy = policy
self.max_layer = max_layer
# For Experiment #
self.random = False
self.gcn = False # GCN Baseline
self.enable_skh = True # only when GCN is false will be useful
self.enable_dlayer = True
self.baseline_experience = 50
# buffers for updating
#self.buffers = {i: [] for i in range(max_layer)}
self.buffers = defaultdict(list)
self.past_performance = [0]
def test(model, loader, loss, device, log_dir, epoch):
model.eval()
error = 0
for i, data in tqdm(enumerate(loader)):
data = data.to(device)
out, mask = model(data)
result = loss(out, to_dense_adj(data.edge_index, data.batch,
data.y)) / mask.sum()
error += result.item()
if i == 0:
if epoch < 5:
check_ts(data, log_dir, epoch)
# divides by number of molecules
return math.sqrt(error / len(loader.dataset)) # rmse
def forward(self, data):
seq_len = data['s']
inputs = data['c'].reshape((len(seq_len), -1, 3))
inputs = inputs.reshape((len(seq_len), -1, 3))
_, idx_sort = torch.sort(seq_len, dim=0, descending=True)
_, idx_unsort = torch.sort(idx_sort, dim=0)
input_x = inputs.index_select(0, Variable(idx_sort))
length_list = list(seq_len[idx_sort])
input_x = input_x.float()
pack = nn.utils.rnn.pack_padded_sequence(input_x,
length_list,
batch_first=True)
out, state = self.lstm(pack)
del state
un_padded = nn.utils.rnn.pad_packed_sequence(out, batch_first=True)
un_padded = un_padded[0].index_select(0, Variable(idx_unsort))
out = self.dropout(un_padded)
feature = self.fc(out)
batch_feature = None
del out, pack, un_padded
for i in range(data.num_graphs):
emptyfeature = torch.zeros((1, self.fea_dim)).to(device)
fea = torch.cat((feature[i][:(seq_len[i])], emptyfeature))
if batch_feature is None:
batch_feature = fea
else:
batch_feature = torch.cat((batch_feature, fea))
data['x'] = batch_feature
x, edge_index = data.x, data.edge_index
dense_x = utils.to_dense_batch(x, batch=data.batch)
x = dense_x[0]
adj = utils.to_dense_adj(data.edge_index, batch=data.batch)
s = self.gnn1_pool(x, adj)
x = self.gnn1_embed(x, adj)
x, adj, l1, e1 = dense_diff_pool(x, adj, s)
x = self.gnn3_embed(x, adj)
x = x.mean(dim=1)
x1 = self.lin1(x)
x = F.relu(x1)
x = self.lin2(x)
return F.log_softmax(x, dim=-1), x1, l1, e1
def rwr_filter(data, c=0.1):
adj = to_dense_adj(data.edge_index).squeeze()
adj = 0.5 * (adj + adj.T)
adj = adj + torch.eye(adj.shape[0])
d = torch.diag(torch.sum(adj, 1))
d_inv = d**(-0.5)
d_inv[torch.isinf(d_inv)] = 0.
w_tilda = torch.matmul(d_inv, adj)
w_tilda = np.matmul(w_tilda, d_inv)
q = torch.eye(w_tilda.shape[0]) - c * w_tilda
q_inv = torch.inverse(q)
rwr = (1 - c) * q_inv
rwr, _ = torch.sort(rwr, dim=1, descending=True)
sparse_rwr = rwr.to_sparse()
data.edge_index = sparse_rwr.indices()
data.edge_attr = sparse_rwr.values().unsqueeze(1).float()
return data
def convert_to_dense(data, mask, start=0, end=None):
end = data.T if not end else end
adjs = []
for t in range(start, end):
ei = data.get_masked_edges(t, mask)
ei = add_remaining_self_loops(ei, num_nodes=data.num_nodes)[0]
a = to_dense_adj(ei, max_num_nodes=data.num_nodes)[0]
d = a.sum(dim=1)
d = 1/torch.sqrt(d)
d = torch.diag(d)
ahat = d @ a @ d
adjs.append(ahat)
return adjs, [torch.eye(data.num_nodes) for _ in range(len(adjs))]
def test_to_dense_adj():
edge_index = torch.tensor([
[0, 0, 1, 2, 3, 4],
[0, 1, 0, 3, 4, 2],
])
batch = torch.tensor([0, 0, 1, 1, 1])
adj = to_dense_adj(edge_index, batch)
assert adj.size() == (2, 3, 3)
assert adj[0].tolist() == [[1, 1, 0], [1, 0, 0], [0, 0, 0]]
assert adj[1].tolist() == [[0, 1, 0], [0, 0, 1], [1, 0, 0]]
adj = to_dense_adj(edge_index, batch, max_num_nodes=5)
assert adj.size() == (2, 5, 5)
assert adj[0][:3, :3].tolist() == [[1, 1, 0], [1, 0, 0], [0, 0, 0]]
assert adj[1][:3, :3].tolist() == [[0, 1, 0], [0, 0, 1], [1, 0, 0]]
edge_attr = torch.Tensor([1, 2, 3, 4, 5, 6])
adj = to_dense_adj(edge_index, batch, edge_attr)
assert adj.size() == (2, 3, 3)
assert adj[0].tolist() == [[1, 2, 0], [3, 0, 0], [0, 0, 0]]
assert adj[1].tolist() == [[0, 4, 0], [0, 0, 5], [6, 0, 0]]
adj = to_dense_adj(edge_index, batch, edge_attr, max_num_nodes=5)
assert adj.size() == (2, 5, 5)
assert adj[0][:3, :3].tolist() == [[1, 2, 0], [3, 0, 0], [0, 0, 0]]
assert adj[1][:3, :3].tolist() == [[0, 4, 0], [0, 0, 5], [6, 0, 0]]
edge_attr = edge_attr.view(-1, 1)
adj = to_dense_adj(edge_index, batch, edge_attr)
assert adj.size() == (2, 3, 3, 1)
edge_attr = edge_attr.view(-1, 1)
adj = to_dense_adj(edge_index, batch, edge_attr, max_num_nodes=5)
assert adj.size() == (2, 5, 5, 1)
adj = to_dense_adj(edge_index)
assert adj.size() == (1, 5, 5)
assert adj[0].nonzero(as_tuple=False).t().tolist() == edge_index.tolist()
adj = to_dense_adj(edge_index, max_num_nodes=10)
assert adj.size() == (1, 10, 10)
assert adj[0].nonzero(as_tuple=False).t().tolist() == edge_index.tolist()
def get_edge_acc(pred_adj, target_data):
# pred_adj.cpu()
# pred_adj = torch.pow(pred_adj,2)
# print(pred_adj.shape)
pred_adj[pred_adj >= args.threshold] = 1
pred_adj[pred_adj < args.threshold] = 0
edge_index, _ = add_self_loops(target_data['edge_index'])
target_adj = to_dense_adj(edge_index)[0]
tn, fp, fn, tp = confusion_matrix(
target_adj.view(-1).cpu().detach().numpy(),
pred_adj.view(-1).cpu().detach().numpy()).ravel()
# print(target_adj.shape)
# check_adj = pred_adj + target_adj
# check_adj[check_adj >= 1] = 1
# check_adj[check_adj < 1] = 0
return (tn + tp) / (tn + fp + fn + tp), tn, fp, fn, tp
def forward(self, x, edge_index, batch):
x = self.conv1(x, edge_index).relu()
x, mask = to_dense_batch(x, batch)
adj = to_dense_adj(edge_index, batch)
_, x, adj, sp1, o1, c1 = self.pool1(x, adj, mask)
x = self.conv2(x, adj).relu()
_, x, adj, sp2, o2, c2 = self.pool2(x, adj)
x = self.conv3(x, adj)
x = x.mean(dim=1)
x = self.lin1(x).relu()
x = self.lin2(x)
return F.log_softmax(x, dim=-1), sp1 + sp2 + o1 + o2 + c1 + c2
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
x, mask = to_dense_batch(x, batch=batch)
adj = to_dense_adj(edge_index, batch=batch)
# data = ToDense(data.num_nodes)(data)
# TODO describe mask shape and how batching works
# adj, mask, x = data.adj, data.mask, data.x
x_all, l_total, e_total = [], 0, 0
for i in range(self.num_diffpool_layers):
if i != 0:
mask = None
x, adj, l, e = self.diffpool_layers[i](
x, adj,
mask) # x has shape (batch, MAX_no_nodes, feature_size)
x_all.append(torch.max(x, dim=1)[0])
l_total += l
e_total += e
x = self.final_embed(x, adj)
x_all.append(torch.max(x, dim=1)[0])
x = torch.cat(x_all,
dim=1) # shape (batch, feature_size x diffpool layers)
x = F.relu(self.lin1(x))
x = self.lin2(x)
# Don't apply sigmoid during training b/c using BCEWithLogitsLoss
if self.classification and not self.training:
x = self.sigmoid(x)
if self.multiclass:
x = x.reshape(
(x.size(0), -1, self.multiclass_num_classes
)) # batch size x num targets x num classes per target
if not self.training:
x = self.multiclass_softmax(
x
) # to get probabilities during evaluation, but not during training as we're using CrossEntropyLoss
return x, l_total, e_total
def forward(self, x, adj, batch_num_nodes=None, **kwargs):
# mask
# Convert ajd matrix
adj = to_dense_adj(adj, max_num_nodes=x.shape[0])
max_num_nodes = adj.size()[1]
if batch_num_nodes is not None:
embedding_mask = self.construct_mask(max_num_nodes,
batch_num_nodes)
else:
embedding_mask = None
self.adj_atts = []
self.embedding_tensor, adj_att = self.gcn_forward(
x, adj, self.conv_first, self.conv_block, self.conv_last,
embedding_mask)
pred = self.pred_model(self.embedding_tensor)
pred = pred.squeeze(0)
return F.log_softmax(pred, dim=1) # pred
def __call__(self, data):
x = data.x if data.x.dim() == 1 else data.x[:, 0]
adj = to_dense_adj(data.edge_index,
data.batch,
edge_attr=data.edge_attr,
max_num_nodes=self.max_size)
x_node, x_node_sizes = to_dense_batch(x,
data.batch,
max_num_nodes=self.max_size)
data.num_atoms = x_node_sizes
data.dense_adj = _one_hotify_adj(adj).float()
try:
data.dense_x = torch.nn.functional.one_hot(
x_node, self.num_atom_classes).float()
except: # noqa: E722 # Ignore bare except since the error is raised anyway.
print(self.num_atom_classes, ", but saw: ", torch.max(x_node))
raise
return data
def generate_negative_samples(data, num_samples):
# Build inverse adj matrix to find non-edges
a = to_dense_adj(data.edge_index[:, data.train_mask])[-1].bool()
# Determine what nodes we'll have embeddings for
known = data.edge_index[:, data.train_mask].flatten().unique()
not_known = torch.full((a.size()[0],), 1, dtype=torch.bool)
not_known[known] = 0
# Mark all untrained nodes as having edges so they won't show
# up in the inverted adj matrix
a[not_known, :] = 1
a[:, not_known] = 1
# That way, when we invert the matrix, the non-edges will only be
# between nodes that we for sure have embeddings for
non_edges = (~a).nonzero()
# Randomly sample from the allowed edge pool
return non_edges[torch.randperm(non_edges.size()[0])[:num_samples], :]
def forward(self, x, edge_index, batch):
x = F.relu(self.conv1(x, edge_index))
x, mask = to_dense_batch(x, batch)
adj = to_dense_adj(edge_index, batch)
s = self.pool1(x)
x, adj, mc1, o1 = dense_mincut_pool(x, adj, s, mask)
x = F.relu(self.conv2(x, adj))
s = self.pool2(x)
x, adj, mc2, o2 = dense_mincut_pool(x, adj, s)
x = self.conv3(x, adj)
x = x.mean(dim=1)
x = F.relu(self.lin1(x))
x = self.lin2(x)
return F.log_softmax(x, dim=-1), mc1 + mc2, o1 + o2
def test(self, loader):
test_acc, _ = self.eval(loader)
exact_match = 0
validity = 0
num_valid = 0
num_invalid = 0
for _, data in enumerate(tqdm(loader, desc='[Test]')):
nodes = data.x.argmax(-1)
a = to_dense_adj(data.edge_index, edge_attr=data.edge_attr)[0]
mol = mol_from_graphs(nodes, a)
if mol is None:
num_invalid += 1
continue
else:
num_valid += 1
smiles = rdkit.Chem.MolToSmiles(mol)
data = data.to(self.args.device)
out = self.model(data)
pred = torch.softmax(out, dim=-1).argmax(-1)
pred_mol = mol_from_graphs(pred, a)
if pred_mol is not None:
pred_smiles = rdkit.Chem.MolToSmiles(pred_mol)
else:
pred_smiles = ''
_exact_match = smiles == pred_smiles
_validity = pred_smiles != ''
exact_match += _exact_match
validity += _validity
exact_match = exact_match / float(num_valid) * 100
validity = validity / float(num_valid) * 100
return num_valid, num_invalid, exact_match, validity, test_acc
def do_trans(data):
node_num, _ = data.x.size()
sub_num = int(node_num * ratio)
if add_self_loop:
sl = torch.tensor([[n, n] for n in range(node_num)]).t()
edge_index = torch.cat((data.edge_index, sl), dim=1)
else:
edge_index = data.edge_index.detach().clone()
# edge_index = edge_index.numpy()
idx_sub = [np.random.randint(node_num, size=1)[0]]
# idx_neigh = set([n for n in edge_index[1][edge_index[0]==idx_sub[0]]])
idx_neigh = set(
[n.item() for n in edge_index[1][edge_index[0] == idx_sub[0]]])
count = 0
while len(idx_sub) <= sub_num:
count = count + 1
if count > node_num:
break
if len(idx_neigh) == 0:
break
sample_node = np.random.choice(list(idx_neigh))
if sample_node in idx_sub:
continue
idx_sub.append(sample_node)
# idx_neigh.union(set([n for n in edge_index[1][edge_index[0]==idx_sub[-1]]]))
idx_neigh.union(
set([
n.item()
for n in edge_index[1][edge_index[0] == idx_sub[-1]]
]))
idx_drop = [n for n in range(node_num) if not n in idx_sub]
idx_sampled = idx_sub
adj = to_dense_adj(edge_index)[0]
adj = adj[idx_sampled, :][:, idx_sampled]
return Data(x=data.x[idx_sampled], edge_index=dense_to_sparse(adj)[0])
def forward(self, obs, return_attention_weights=False, **kwargs):
batch_size = obs.shape[0]
x, edge_index = self.pre_graph_builder(obs)
x, edge_attention = self.attentioner(x, edge_index, return_attention_weights=True)
x = self.hidden_activations[0](x)
edge_index, attention_weights = edge_attention
attention_weights = attention_weights.squeeze(-1)
batch_vector = torch.arange(batch_size)[:,None].repeat(1,self.pre_graph_builder.node_num).reshape(-1)
adj = pyg_utils.to_dense_adj(edge_index, edge_attr=attention_weights, batch=batch_vector)
adj = adj.squeeze(-1)
x, _ = pyg_utils.to_dense_batch(x, batch=batch_vector)
for l, conv in enumerate(self.convs):
x = conv(x, adj)
x = self.hidden_activations[l+1](x)
# for l, conv in enumerate(self.convs):
# x = conv(x, edge_index, edge_weight = attention_weights)
# x = self.hidden_activations[l+1](x)
x = x.reshape(batch_size,self.pre_graph_builder.node_num,self.node_dim)
x = self.output_activation(x)
if return_attention_weights:
return x, edge_attention
else:
return x
def test_dense_graph_conv(aggr):
channels = 16
sparse_conv = GraphConv(channels, channels, aggr=aggr)
dense_conv = DenseGraphConv(channels, channels, aggr=aggr)
assert dense_conv.__repr__() == 'DenseGraphConv(16, 16)'
# Ensure same weights and bias.
dense_conv.lin_rel = sparse_conv.lin_rel
dense_conv.lin_root = sparse_conv.lin_root
x = torch.randn((5, channels))
edge_index = torch.tensor([[0, 0, 1, 1, 2, 2, 3, 4],
[1, 2, 0, 2, 0, 1, 4, 3]])
sparse_out = sparse_conv(x, edge_index)
assert sparse_out.size() == (5, channels)
adj = to_dense_adj(edge_index)
mask = torch.ones(5, dtype=torch.bool)
dense_out = dense_conv(x, adj, mask)[0]
assert dense_out.size() == (5, channels)
assert torch.allclose(sparse_out, dense_out, atol=1e-04)
def train(model, loader, loss_func, opt):
total_loss = 0
model.train()
for batch_id, rxn_batch in enumerate(tqdm(loader)):
opt.zero_grad()
rxn_batch = rxn_batch.to(model.device)
D_pred, mask, _ = model(rxn_batch)
D_gt = to_dense_adj(rxn_batch.edge_index, rxn_batch.batch, rxn_batch.y)
batch_loss = loss_func(D_pred, D_gt) / mask.sum()
batch_loss.backward()
clip_grad_norm_(model.parameters(), MAX_CLIP_NORM) # clip gradients
# if logger:
# pnorm = compute_parameters_norm(model)
# gnorm = compute_gradients_norm(model)
# logger.info(f' Batch {batch_id} Loss: {batch_loss.item()}\t Parameter Norm: {pnorm}\t Gradient Norm: {gnorm}')
opt.step()
total_loss += batch_loss.item()
RMSE = math.sqrt(total_loss / len(loader.dataset))
return RMSE
def pg_graph_to_nb201(pg_graph):
# first tensor node attributes, second is the edge list
ops = [OPS_by_IDX_201[i] for i in pg_graph.x.cpu().numpy()]
matrix = np.array(to_dense_adj(pg_graph.edge_index)[0].cpu().numpy())
try:
if (matrix == ADJACENCY).all():
steps_coding = ['0', '0', '1', '0', '1', '2']
node_1 = '|' + ops[1] + '~' + steps_coding[0] + '|'
node_2 = '|' + ops[2] + '~' + steps_coding[1] + '|' + ops[
3] + '~' + steps_coding[2] + '|'
node_3 = '|' + ops[4] + '~' + steps_coding[3] + '|' + ops[
5] + '~' + steps_coding[4] + '|' + ops[
6] + '~' + steps_coding[5] + '|'
nodes_nb201 = node_1 + '+' + node_2 + '+' + node_3
index = nasbench.query_index_by_arch(nodes_nb201)
acc = Dataset.map_item(index).acc
else:
acc = torch.zeros(1)
except:
acc = torch.zeros(1)
return acc
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
if self.encode_edge:
x = self.atom_encoder(x)
x = F.relu(self.convs[0](x, edge_index, data.edge_attr))
else:
x = F.relu(self.convs[0](x, edge_index))
x, mask = to_dense_batch(x, batch)
adj = to_dense_adj(edge_index, batch)
if self.pooling_type != 'mlp':
s = self.rms[0][:x.size(1), :].unsqueeze(dim=0).expand(
x.size(0), -1, -1).to(x.device)
else:
s = self.pools[0](x)
x, adj, mc, o = dense_mincut_pool(x, adj, s, mask)
for i in range(1, self.num_layers - 1):
x = F.relu(self.convs[i](x, adj))
if self.pooling_type != 'mlp':
s = self.rms[i][:x.size(1), :].unsqueeze(dim=0).expand(
x.size(0), -1, -1).to(x.device)
else:
s = self.pools[i](x)
x, adj, mc_aux, o_aux = dense_mincut_pool(x, adj, s)
mc += mc_aux
o += o_aux
x = self.convs[self.num_layers - 1](x, adj)
x = x.mean(dim=1)
x = F.relu(self.lin1(x))
x = self.lin2(x)
return x, mc, o
def dynamic_new_link_prediction(data,
partition_fn,
zs,
start=0,
end=None,
include_tr=True,
batched=False):
if batched:
raise NotImplementedError("Sorry, batching is a TODO")
p, n = [], []
b = None
if partition_fn == None:
partition_fn = lambda x: data.eis[x]
end = end if end else data.T
for i in range(start, end):
# Use full adj matrix for new link pred
ei = partition_fn(i)
a = b
b = to_dense_adj(ei, max_num_nodes=data.num_nodes)[0].bool()
if type(a) == type(None):
continue
# Generates new links in next time step
new_links = (~a).logical_and(a.logical_or(b))
new_links, _ = dense_to_sparse(new_links)
p.append(new_links)
n.append(fast_negative_sampling(ei, p[-1].size(1), data.num_nodes))
return p, n, zs[:-1]
