iicd/Graph-DiT
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

use trainer but has bugs

Browse Source
This commit is contained in:
mhz
2024-09-19 14:11:19 +02:00
parent d36e1d1077
commit be178bc5ee
6 changed files with 750 additions and 580 deletions
Download Patch File Download Diff File Expand all files Collapse all files

135
graph_dit/diffusion_model.py
View File

@@ -23,6 +23,9 @@ class Graph_DiT(pl.LightningModule):
self.test_only = cfg.general.test_only
self.guidance_target = getattr(cfg.dataset, 'guidance_target', None)
from nas_201_api import NASBench201API as API
self.api = API('/nfs/data3/hanzhang/nasbenchDiT/graph_dit/NAS-Bench-201-v1_1-096897.pth')
input_dims = dataset_infos.input_dims
output_dims = dataset_infos.output_dims
nodes_dist = dataset_infos.nodes_dist
@@ -79,6 +82,7 @@ class Graph_DiT(pl.LightningModule):
self.node_dist = nodes_dist
self.active_index = active_index
self.dataset_info = dataset_infos
self.cur_epoch = 0
self.train_loss = TrainLossDiscrete(self.cfg.model.lambda_train)
@@ -162,25 +166,81 @@ class Graph_DiT(pl.LightningModule):
return pred
def training_step(self, data, i):
data_x = F.one_hot(data.x, num_classes=12).float()[:, self.active_index]
data_edge_attr = F.one_hot(data.edge_attr, num_classes=2).float()
if self.cfg.general.type != 'accelerator' and self.current_epoch > self.cfg.train.n_epochs / 5 * 4:
samples_left_to_generate = self.cfg.general.samples_to_generate
samples_left_to_save = self.cfg.general.samples_to_save
chains_left_to_save = self.cfg.general.chains_to_save
dense_data, node_mask = utils.to_dense(data_x, data.edge_index, data_edge_attr, data.batch, self.max_n_nodes)
dense_data = dense_data.mask(node_mask)
X, E = dense_data.X, dense_data.E
noisy_data = self.apply_noise(X, E, data.y, node_mask)
pred = self.forward(noisy_data)
loss = self.train_loss(masked_pred_X=pred.X, masked_pred_E=pred.E, pred_y=pred.y,
true_X=X, true_E=E, true_y=data.y, node_mask=node_mask,
samples, all_ys, batch_id = [], [], 0
def graph_reward_fn(graphs, true_graphs=None, device=None, reward_model='swap'):
rewards = []
if reward_model == 'swap':
import csv
with open('/nfs/data3/hanzhang/nasbenchDiT/graph_dit/swap_results.csv', 'r') as f:
reader = csv.reader(f)
header = next(reader)
data = [row for row in reader]
swap_scores = [float(row[0]) for row in data]
for graph in graphs:
node_tensor = graph[0]
node = node_tensor.cpu().numpy().tolist()
def nodes_to_arch_str(nodes):
num_to_op = ['input', 'nor_conv_1x1', 'nor_conv_3x3', 'avg_pool_3x3', 'skip_connect', 'none', 'output']
nodes_str = [num_to_op[node] for node in nodes]
arch_str = '|' + nodes_str[1] + '~0|+' + \
'|' + nodes_str[2] + '~0|' + nodes_str[3] + '~1|+' +\
'|' + nodes_str[4] + '~0|' + nodes_str[5] + '~1|' + nodes_str[6] + '~2|'
return arch_str
arch_str = nodes_to_arch_str(node)
reward = swap_scores[self.api.query_index_by_arch(arch_str)]
rewards.append(reward)
return torch.tensor(rewards, dtype=torch.float32, requires_grad=True).unsqueeze(0).to(device)
old_log_probs = None
bs = 1 * self.cfg.train.batch_size
to_generate = min(samples_left_to_generate, bs)
to_save = min(samples_left_to_save, bs)
chains_save = min(chains_left_to_save, bs)
# batch_y = test_y_collection[batch_id : batch_id + to_generate]
batch_y = torch.ones(to_generate, self.ydim_output, device=self.device)
cur_sample, log_probs = self.sample_batch(batch_id, to_generate, batch_y, save_final=to_save,
keep_chain=chains_save, number_chain_steps=self.number_chain_steps)
# samples = samples + cur_sample
samples.append(cur_sample)
reward = graph_reward_fn(cur_sample, device=self.device)
advantages = (reward - torch.mean(reward)) / (torch.std(reward) + 1e-6) #
if old_log_probs is None:
old_log_probs = log_probs.clone()
ratio = torch.exp(log_probs - old_log_probs)
print(f"ratio: {ratio.shape}, advantages: {advantages.shape}")
unclipped_loss = -advantages * ratio
clipped_loss = -advantages * torch.clamp(ratio, 1.0 - self.cfg.ppo.clip_param, 1.0 + self.cfg.ppo.clip_param)
loss = torch.mean(torch.max(unclipped_loss, clipped_loss))
return {'loss': loss}
else:
data_x = F.one_hot(data.x, num_classes=12).float()[:, self.active_index]
data_edge_attr = F.one_hot(data.edge_attr, num_classes=2).float()
dense_data, node_mask = utils.to_dense(data_x, data.edge_index, data_edge_attr, data.batch, self.max_n_nodes)
dense_data = dense_data.mask(node_mask)
X, E = dense_data.X, dense_data.E
noisy_data = self.apply_noise(X, E, data.y, node_mask)
pred = self.forward(noisy_data)
loss = self.train_loss(masked_pred_X=pred.X, masked_pred_E=pred.E, pred_y=pred.y,
true_X=X, true_E=E, true_y=data.y, node_mask=node_mask,
log=i % self.log_every_steps == 0)
# print(f'training loss: {loss}, epoch: {self.current_epoch}, batch: {i}\n, pred type: {type(pred)}, pred.X shape: {type(pred.X)}, {pred.X.shape}, pred.E shape: {type(pred.E)}, {pred.E.shape}')
self.train_metrics(masked_pred_X=pred.X, masked_pred_E=pred.E, true_X=X, true_E=E,
log=i % self.log_every_steps == 0)
# print(f'training loss: {loss}, epoch: {self.current_epoch}, batch: {i}\n, pred type: {type(pred)}, pred.X shape: {type(pred.X)}, {pred.X.shape}, pred.E shape: {type(pred.E)}, {pred.E.shape}')
self.train_metrics(masked_pred_X=pred.X, masked_pred_E=pred.E, true_X=X, true_E=E,
log=i % self.log_every_steps == 0)
self.log(f'loss', loss, batch_size=X.size(0), sync_dist=True)
print(f"training loss: {loss}")
with open("training-loss.csv", "a") as f:
f.write(f"{loss}, {i}\n")
return {'loss': loss}
self.log(f'loss', loss, batch_size=X.size(0), sync_dist=True)
print(f"training loss: {loss}")
with open("training-loss.csv", "a") as f:
f.write(f"{loss}, {i}\n")
return {'loss': loss}
def configure_optimizers(self):
@@ -196,14 +256,15 @@ class Graph_DiT(pl.LightningModule):
def on_train_epoch_start(self) -> None:
if self.current_epoch / self.trainer.max_epochs in [0.25, 0.5, 0.75, 1.0]:
print("Starting train epoch {}/{}...".format(self.current_epoch, self.trainer.max_epochs))
# if self.cur_epoch / self.cfg.train.n_epochs in [0.25, 0.5, 0.75, 1.0]:
print("Starting train epoch {}/{}...".format(self.cur_epoch, self.cfg.train.n_epochs))
self.start_epoch_time = time.time()
self.train_loss.reset()
self.train_metrics.reset()
def on_train_epoch_end(self) -> None:
if self.current_epoch / self.trainer.max_epochs in [0.25, 0.5, 0.75, 1.0]:
if self.current_epoch / self.cfg.train.n_epochs in [0.25, 0.5, 0.75, 1.0]:
log = True
else:
log = False
@@ -240,6 +301,7 @@ class Graph_DiT(pl.LightningModule):
self.val_X_logp.compute(), self.val_E_logp.compute()]
if self.current_epoch / self.trainer.max_epochs in [0.25, 0.5, 0.75, 1.0]:
# if self.cur_epoch / self.cfg.train.n_epochs in [0.25, 0.5, 0.75, 1.0]:
print(f"Epoch {self.current_epoch}: Val NLL {metrics[0] :.2f} -- Val Atom type KL {metrics[1] :.2f} -- ",
f"Val Edge type KL: {metrics[2] :.2f}", 'Val loss: %.2f \t Best : %.2f\n' % (metrics[0], self.best_val_nll))
with open("validation-metrics.csv", "a") as f:
@@ -336,7 +398,7 @@ class Graph_DiT(pl.LightningModule):
print(f"Epoch {self.current_epoch}: Test NLL {metrics[0] :.2f} -- Test Atom type KL {metrics[1] :.2f} -- ",
f"Test Edge type KL: {metrics[2] :.2f}")
## final epcoh
## final epoch
samples_left_to_generate = self.cfg.general.final_model_samples_to_generate
samples_left_to_save = self.cfg.general.final_model_samples_to_save
chains_left_to_save = self.cfg.general.final_model_chains_to_save
@@ -359,9 +421,9 @@ class Graph_DiT(pl.LightningModule):
# batch_y = test_y_collection[batch_id : batch_id + to_generate]
batch_y = torch.ones(to_generate, self.ydim_output, device=self.device)
cur_sample = self.sample_batch(batch_id, to_generate, batch_y, save_final=to_save,
cur_sample, log_probs = self.sample_batch(batch_id, to_generate, batch_y, save_final=to_save,
keep_chain=chains_save, number_chain_steps=self.number_chain_steps)
samples = samples + cur_sample
samples.append(cur_sample)
all_ys.append(batch_y)
batch_id += to_generate
@@ -601,6 +663,12 @@ class Graph_DiT(pl.LightningModule):
assert (E == torch.transpose(E, 1, 2)).all()
if self.cfg.general.type != 'accelerator':
if self.trainer.training or self.trainer.validating:
total_log_probs = torch.zeros([self.cfg.general.samples_to_generate, 10], device=self.device)
elif self.trainer.testing:
total_log_probs = torch.zeros([self.cfg.general.final_model_samples_to_generate, 10], device=self.device)
# Iteratively sample p(z_s | z_t) for t = 1, ..., T, with s = t - 1.
for s_int in reversed(range(0, self.T)):
s_array = s_int * torch.ones((batch_size, 1)).type_as(y)
@@ -609,21 +677,24 @@ class Graph_DiT(pl.LightningModule):
t_norm = t_array / self.T
# Sample z_s
sampled_s, discrete_sampled_s = self.sample_p_zs_given_zt(s_norm, t_norm, X, E, y, node_mask)
sampled_s, discrete_sampled_s, log_probs = self.sample_p_zs_given_zt(s_norm, t_norm, X, E, y, node_mask)
X, E, y = sampled_s.X, sampled_s.E, sampled_s.y
total_log_probs += log_probs
# Sample
sampled_s = sampled_s.mask(node_mask, collapse=True)
X, E, y = sampled_s.X, sampled_s.E, sampled_s.y
molecule_list = []
graph_list = []
for i in range(batch_size):
n = n_nodes[i]
atom_types = X[i, :n].cpu()
node_types = X[i, :n].cpu()
edge_types = E[i, :n, :n].cpu()
molecule_list.append([atom_types, edge_types])
graph_list.append((node_types , edge_types))
return molecule_list
total_log_probs = torch.sum(total_log_probs, dim=-1)
return graph_list, total_log_probs
def sample_p_zs_given_zt(self, s, t, X_t, E_t, y_t, node_mask):
"""Samples from zs ~ p(zs | zt). Only used during sampling.
@@ -675,6 +746,14 @@ class Graph_DiT(pl.LightningModule):
# with condition = P_t(A_{t-1} |A_t, y)
prob_X, prob_E, pred = get_prob(noisy_data)
log_prob_X = torch.log(torch.gather(prob_X, -1, X_t.long()).squeeze(-1)) # bs, n
log_prob_E = torch.log(torch.gather(prob_E, -1, E_t.long()).squeeze(-1)) # bs, n, n
# Sum the log_prob across dimensions for total log_prob
log_prob_X = log_prob_X.sum(dim=-1)
log_prob_E = log_prob_E.sum(dim=(1, 2))
log_probs = torch.cat([log_prob_X, log_prob_E], dim=-1)
### Guidance
if self.guidance_target is not None and self.guide_scale is not None and self.guide_scale != 1:
uncon_prob_X, uncon_prob_E, pred = get_prob(noisy_data, unconditioned=True)
@@ -810,4 +889,4 @@ class Graph_DiT(pl.LightningModule):
out_one_hot = utils.PlaceHolder(X=X_s, E=E_s, y=y_t)
out_discrete = utils.PlaceHolder(X=X_s, E=E_s, y=y_t)
return out_one_hot.mask(node_mask).type_as(y_t), out_discrete.mask(node_mask, collapse=True).type_as(y_t)
return out_one_hot.mask(node_mask).type_as(y_t), out_discrete.mask(node_mask, collapse=True).type_as(y_t), log_probs