def __init__(self): # Can manually select the device too self.device = torch.device( 'cuda:0' if torch.cuda.is_available() else 'cpu') self.model_path = 'model_save/' + 'model_name' self.num_epochs = get_model_attribute( 'epoch', self.model_path, self.device) # Whether to generate networkx format graphs for real datasets self.generate_graphs = True self.count = 2560 self.batch_size = 32 # Must be a factor of count self.metric_eval_batch_size = 256 # Specific DFScodeRNN self.max_num_edges = 50 # Specific to GraphRNN self.min_num_node = 0 self.max_num_node = 40 self.train_args = get_model_attribute( 'saved_args', self.model_path, self.device) self.graphs_save_path = 'graphs/' self.current_graphs_save_path = self.graphs_save_path + self.train_args.fname + '_' + \ self.train_args.time + '/' + str(self.num_epochs) + '/'
def train(args, dataloader_train, model, feature_map, dataloader_validate=None): # initialize optimizer optimizer = {} for name, net in model.items(): optimizer['optimizer_' + name] = optim.Adam( list(net.parameters()), lr=args.lr, weight_decay=5e-5) scheduler = {} for name, net in model.items(): optimizer['scheduler_' + name] = MultiStepLR( optimizer['optimizer_' + name], milestones=args.milestones, gamma=args.gamma) if args.load_model: load_model(args.load_model_path, args.device, model, optimizer, scheduler) print('Model loaded') epoch = get_model_attribute('epoch', args.load_model_path, args.device) else: epoch = 0 while epoch < args.epochs: loss = train_epoch(epoch, args, model, dataloader_train, optimizer, scheduler, feature_map) epoch += 1 # logging if args.log_tensorboard: log_value('train_loss ' + args.fname, loss, epoch) else: print('Epoch: {}/{}, train loss: {:.6f}'.format(epoch, args.epochs, loss)) # save model checkpoint if args.save_model and epoch != 0 and epoch % args.epochs_save == 0: save_model(epoch, args, model, optimizer, scheduler, feature_map=feature_map) print('Model Saved - Epoch: {}/{}, train loss: {:.6f}'.format(epoch, args.epochs, loss)) if dataloader_validate is not None and epoch % args.epochs_validate == 0: loss_validate = test_data(args, model, dataloader_validate, feature_map) if args.log_tensorboard: log_value('validate_loss ' + args.fname, loss_validate, epoch) else: print('Epoch: {}/{}, validation loss: {:.6f}'.format(epoch, args.epochs, loss_validate)) save_model(epoch, args, model, optimizer, scheduler, feature_map=feature_map) print('Model Saved - Epoch: {}/{}, train loss: {:.6f}'.format(epoch, args.epochs, loss))
def update_args(self): if self.load_model: args = get_model_attribute( 'saved_args', self.load_model_path, self.load_device) args.device = self.load_device args.load_model = True args.load_model_path = self.load_model_path args.epochs = self.epochs_end args.clean_tensorboard = False args.clean_temp = False args.produce_graphs = False args.produce_min_dfscodes = False args.produce_min_dfscode_tensors = False return args return self
def predict_graphs(eval_args): train_args = eval_args.train_args feature_map = get_model_attribute('feature_map', eval_args.model_path, eval_args.device) train_args.device = eval_args.device model = create_model(train_args, feature_map) load_model(eval_args.model_path, eval_args.device, model) for _, net in model.items(): net.eval() graphs = [] for _ in range(eval_args.count // eval_args.batch_size): sampled_graphs = model['generator'](eval_args.batch_size, training=False) nb = feature_map['node_backward'] eb = feature_map['edge_backward'] for sampled_graph in sampled_graphs: graph = sampled_graph.to_networkx(node_attrs=['label'], edge_attrs=['label' ]).to_undirected() labeled_graph = nx.Graph() for v in graph.nodes(): labeled_graph.add_node( v, label=nb[graph.nodes[v]['label'].item() - 1]) for u, v in graph.edges(): labeled_graph.add_edge( u, v, label=eb[graph.edges[u, v]['label'].item()]) # Take maximum connected component if len(labeled_graph.nodes()) > 0: max_comp = max(nx.connected_components(labeled_graph), key=len) labeled_graph = labeled_graph.subgraph(max_comp) graphs.append(labeled_graph) return graphs
def train(args, dataloader_train, model, feature_map, dataloader_validate=None): # initialize optimizer optimizer = {} for name, net in model.items(): optimizer['optimizer_' + name] = optim.Adam( filter(lambda p: p.requires_grad, net.parameters()), lr=args.lr, weight_decay=5e-5) scheduler = {} for name, net in model.items(): scheduler['scheduler_' + name] = MultiStepLR( optimizer['optimizer_' + name], milestones=args.milestones, gamma=args.gamma) if args.load_model: load_model(args.load_model_path, args.device, model, optimizer, scheduler) print('Model loaded') epoch = get_model_attribute('epoch', args.load_model_path, args.device) else: epoch = 0 if args.log_tensorboard: writer = SummaryWriter( log_dir=args.tensorboard_path + args.fname + ' ' + args.time, flush_secs=5) else: writer = None while epoch < args.epochs: loss = train_epoch( epoch, args, model, dataloader_train, optimizer, scheduler, feature_map, writer) epoch += 1 # logging if args.log_tensorboard: writer.add_scalar('{} {} Loss/train'.format( args.note, args.graph_type), loss, epoch) else: print('Epoch: {}/{}, train loss: {:.6f}'.format(epoch, args.epochs, loss)) # save model checkpoint if args.save_model and epoch != 0 and epoch % args.epochs_save == 0: save_model( epoch, args, model, optimizer, scheduler, feature_map=feature_map) print( 'Model Saved - Epoch: {}/{}, train loss: {:.6f}'.format(epoch, args.epochs, loss)) if dataloader_validate is not None and epoch % args.epochs_validate == 0: loss_validate = test_data( args, model, dataloader_validate, feature_map) if args.log_tensorboard: writer.add_scalar('{} {} Loss/validate'.format( args.note, args.graph_type), loss_validate, epoch) else: print('Epoch: {}/{}, validation loss: {:.6f}'.format( epoch, args.epochs, loss_validate)) save_model(epoch, args, model, optimizer, scheduler, feature_map=feature_map) print('Model Saved - Epoch: {}/{}, train loss: {:.6f}'.format(epoch, args.epochs, loss))
def predict_graphs(eval_args): train_args = eval_args.train_args feature_map = get_model_attribute('feature_map', eval_args.model_path, eval_args.device) train_args.device = eval_args.device model = create_model(train_args, feature_map) load_model(eval_args.model_path, eval_args.device, model) for _, net in model.items(): net.eval() max_nodes = feature_map['max_nodes'] len_node_vec, len_edge_vec = len(feature_map['node_forward']) + 1, len( feature_map['edge_forward']) + 1 feature_len = 2 * (max_nodes + 1) + 2 * len_node_vec + len_edge_vec graphs = [] for _ in range(eval_args.count // eval_args.batch_size): # initialize dfs_code_rnn hidden according to batch size model['dfs_code_rnn'].hidden = model['dfs_code_rnn'].init_hidden( batch_size=eval_args.batch_size) rnn_input = torch.zeros((eval_args.batch_size, 1, feature_len), device=eval_args.device) pred = torch.zeros((eval_args.batch_size, eval_args.max_num_edges, 5), device=eval_args.device) for i in range(eval_args.max_num_edges): rnn_output = model['dfs_code_rnn'](rnn_input) # Evaluating dfscode tuple timestamp1 = model['output_timestamp1'](rnn_output).reshape( eval_args.batch_size, -1) timestamp2 = model['output_timestamp2'](rnn_output).reshape( eval_args.batch_size, -1) vertex1 = model['output_vertex1'](rnn_output).reshape( eval_args.batch_size, -1) edge = model['output_edge'](rnn_output).reshape( eval_args.batch_size, -1) vertex2 = model['output_vertex2'](rnn_output).reshape( eval_args.batch_size, -1) if train_args.loss_type == 'BCE': timestamp1 = Categorical(timestamp1).sample() timestamp2 = Categorical(timestamp2).sample() vertex1 = Categorical(vertex1).sample() edge = Categorical(edge).sample() vertex2 = Categorical(vertex2).sample() elif train_args.loss_type == 'NLL': timestamp1 = Categorical(logits=timestamp1).sample() timestamp2 = Categorical(logits=timestamp2).sample() vertex1 = Categorical(logits=vertex1).sample() edge = Categorical(logits=edge).sample() vertex2 = Categorical(logits=vertex2).sample() rnn_input = torch.zeros((eval_args.batch_size, 1, feature_len), device=eval_args.device) rnn_input[torch.arange(eval_args.batch_size), 0, timestamp1] = 1 rnn_input[torch.arange(eval_args.batch_size), 0, max_nodes + 1 + timestamp2] = 1 rnn_input[torch.arange(eval_args.batch_size), 0, 2 * max_nodes + 2 + vertex1] = 1 rnn_input[torch.arange(eval_args.batch_size), 0, 2 * max_nodes + 2 + len_node_vec + edge] = 1 rnn_input[torch.arange(eval_args.batch_size), 0, 2 * max_nodes + 2 + len_node_vec + len_edge_vec + vertex2] = 1 pred[:, i, 0] = timestamp1 pred[:, i, 1] = timestamp2 pred[:, i, 2] = vertex1 pred[:, i, 3] = edge pred[:, i, 4] = vertex2 nb = feature_map['node_backward'] eb = feature_map['edge_backward'] for i in range(eval_args.batch_size): dfscode = [] for j in range(eval_args.max_num_edges): if pred[i, j, 0] == max_nodes or pred[i, j, 1] == max_nodes \ or pred[i, j, 2] == len_node_vec - 1 or pred[i, j, 3] == len_edge_vec - 1 \ or pred[i, j, 4] == len_node_vec - 1: break dfscode.append( (int(pred[i, j, 0].data), int(pred[i, j, 1].data), nb[int(pred[i, j, 2].data)], eb[int(pred[i, j, 3].data)], nb[int(pred[i, j, 4].data)])) graph = graph_from_dfscode(dfscode) # Remove self loops graph.remove_edges_from(nx.selfloop_edges(graph)) # Take maximum connected component if len(graph.nodes()): max_comp = max(nx.connected_components(graph), key=len) graph = nx.Graph(graph.subgraph(max_comp)) graphs.append(graph) return graphs
def train(args, dataloader_train, model, feature_map, dataloader_validate=None): # initialize optimizer optimizer = get_optimizer(model, args) scheduler = get_scheduler(model, optimizer, args) if args.load_model: load_model(args.load_model_path, args.device, model, optimizer, scheduler) print('Model loaded') epoch = get_model_attribute('epoch', args.load_model_path, args.device) else: epoch = 0 if args.log_tensorboard: writer = SummaryWriter(log_dir=args.tensorboard_path + args.fname + ' ' + args.time, flush_secs=5) else: writer = None while epoch < args.epochs: loss, acc = train_epoch(epoch, args, model, dataloader_train, optimizer, scheduler, feature_map, writer) epoch += 1 print('Epoch: {}/{}, train loss: {:.3f}, accuray: {:.3f}'.format( epoch, args.epochs, loss, acc)) # logging if args.log_tensorboard: writer.add_scalar( '{} {} Loss/train'.format(args.note, args.graph_type), loss, epoch) # save model checkpoint if args.save_model and epoch != 0 and epoch % args.epochs_save == 0: save_model(epoch, args, model, optimizer, scheduler, feature_map=feature_map) print('Model Saved - Epoch: {}/{}, train loss: {:.6f}'.format( epoch, args.epochs, loss)) if dataloader_validate is not None and epoch % args.epochs_validate == 0: loss_validate = test_data(args, model, dataloader_validate, feature_map) if args.log_tensorboard: writer.add_scalar( '{} {} Loss/validate'.format(args.note, args.graph_type), loss_validate, epoch) else: print('Epoch: {}/{}, validation loss: {:.6f}'.format( epoch, args.epochs, loss_validate)) save_model(epoch, args, model, optimizer, scheduler, feature_map=feature_map) print('Model Saved - Epoch: {}/{}, train loss: {:.6f}'.format( epoch, args.epochs, loss))
def predict_graphs(eval_args): """ Generate graphs (networkx format) given a trained generative graphRNN model :param eval_args: ArgsEvaluate object """ train_args = eval_args.train_args feature_map = get_model_attribute( 'feature_map', eval_args.model_path, eval_args.device) train_args.device = eval_args.device model = create_model(train_args, feature_map) load_model(eval_args.model_path, eval_args.device, model) for _, net in model.items(): net.eval() max_num_node = eval_args.max_num_node len_node_vec, len_edge_vec, num_nodes_to_consider = get_attributes_len_for_graph_rnn( len(feature_map['node_forward']), len(feature_map['edge_forward']), train_args.max_prev_node, train_args.max_head_and_tail) feature_len = len_node_vec + num_nodes_to_consider * len_edge_vec graphs = [] for _ in range(eval_args.count // eval_args.batch_size): model['node_level_rnn'].hidden = model['node_level_rnn'].init_hidden( batch_size=eval_args.batch_size) # [batch_size] * [num of nodes] x_pred_node = np.zeros( (eval_args.batch_size, max_num_node), dtype=np.int32) # [batch_size] * [num of nodes] * [num_nodes_to_consider] x_pred_edge = np.zeros( (eval_args.batch_size, max_num_node, num_nodes_to_consider), dtype=np.int32) node_level_input = torch.zeros( eval_args.batch_size, 1, feature_len, device=eval_args.device) # Initialize to node level start token node_level_input[:, 0, len_node_vec - 2] = 1 for i in range(max_num_node): # [batch_size] * [1] * [hidden_size_node_level_rnn] node_level_output = model['node_level_rnn'](node_level_input) # [batch_size] * [1] * [node_feature_len] node_level_pred = model['output_node'](node_level_output) # [batch_size] * [node_feature_len] for torch.multinomial node_level_pred = node_level_pred.reshape( eval_args.batch_size, len_node_vec) # [batch_size]: Sampling index to set 1 in next node_level_input and x_pred_node # Add a small probability for each node label to avoid zeros node_level_pred[:, :-2] += EPS # Start token should not be sampled. So set it's probability to 0 node_level_pred[:, -2] = 0 # End token should not be sampled if i less than min_num_node if i < eval_args.min_num_node: node_level_pred[:, -1] = 0 sample_node_level_output = torch.multinomial( node_level_pred, 1).reshape(-1) node_level_input = torch.zeros( eval_args.batch_size, 1, feature_len, device=eval_args.device) node_level_input[torch.arange( eval_args.batch_size), 0, sample_node_level_output] = 1 # [batch_size] * [num of nodes] x_pred_node[:, i] = sample_node_level_output.cpu().data # [batch_size] * [1] * [hidden_size_edge_level_rnn] hidden_edge = model['embedding_node_to_edge'](node_level_output) hidden_edge_rem_layers = torch.zeros( train_args.num_layers - 1, eval_args.batch_size, hidden_edge.size(2), device=eval_args.device) # [num_layers] * [batch_size] * [hidden_len] model['edge_level_rnn'].hidden = torch.cat( (hidden_edge.permute(1, 0, 2), hidden_edge_rem_layers), dim=0) # [batch_size] * [1] * [edge_feature_len] edge_level_input = torch.zeros( eval_args.batch_size, 1, len_edge_vec, device=eval_args.device) # Initialize to edge level start token edge_level_input[:, 0, len_edge_vec - 2] = 1 for j in range(min(num_nodes_to_consider, i)): # [batch_size] * [1] * [edge_feature_len] edge_level_output = model['edge_level_rnn'](edge_level_input) # [batch_size] * [edge_feature_len] needed for torch.multinomial edge_level_output = edge_level_output.reshape( eval_args.batch_size, len_edge_vec) # [batch_size]: Sampling index to set 1 in next edge_level input and x_pred_edge # Add a small probability for no edge to avoid zeros edge_level_output[:, -3] += EPS # Start token and end should not be sampled. So set it's probability to 0 edge_level_output[:, -2:] = 0 sample_edge_level_output = torch.multinomial( edge_level_output, 1).reshape(-1) edge_level_input = torch.zeros( eval_args.batch_size, 1, len_edge_vec, device=eval_args.device) edge_level_input[:, 0, sample_edge_level_output] = 1 # Setting edge feature for next node_level_input node_level_input[:, 0, len_node_vec + j * len_edge_vec: len_node_vec + (j + 1) * len_edge_vec] = \ edge_level_input[:, 0, :] # [batch_size] * [num of nodes] * [num_nodes_to_consider] x_pred_edge[:, i, j] = sample_edge_level_output.cpu().data # Save the batch of graphs for k in range(eval_args.batch_size): G = nx.Graph() for v in range(max_num_node): # End node token if x_pred_node[k, v] == len_node_vec - 1: break elif x_pred_node[k, v] < len(feature_map['node_forward']): G.add_node( v, label=feature_map['node_backward'][x_pred_node[k, v]]) else: print('Error in sampling node features') exit() for u in range(len(G.nodes())): for p in range(min(num_nodes_to_consider, u)): if x_pred_edge[k, u, p] < len(feature_map['edge_forward']): if train_args.max_prev_node is not None: v = u - p - 1 elif train_args.max_head_and_tail is not None: if p < train_args.max_head_and_tail[1]: v = u - p - 1 else: v = p - train_args.max_head_and_tail[1] G.add_edge( u, v, label=feature_map['edge_backward'][x_pred_edge[k, u, p]]) elif x_pred_edge[k, u, p] == len(feature_map['edge_forward']): # No edge pass else: print('Error in sampling edge features') exit() # Take maximum connected component if len(G.nodes()): max_comp = max(nx.connected_components(G), key=len) G = nx.Graph(G.subgraph(max_comp)) graphs.append(G) return graphs