def get_train_val_loader(path, train_batch_size=64, val_batch_size=64, val_split=1 / 12): train_dataset = MNISTSuperpixels(path, "train", transform=T.Cartesian()) dataset_size = len(train_dataset) indices = list(range(dataset_size)) split = int(val_split * dataset_size) np.random.seed(43) np.random.shuffle(indices) train_indices = indices[split:] val_indices = indices[:split] train_sampler = SubsetRandomSampler(train_indices) val_sampler = SubsetRandomSampler(val_indices) validation_loader = DataLoader(train_dataset, batch_size=val_batch_size, sampler=val_sampler, shuffle=False) train_loader = DataLoader(train_dataset, batch_size=train_batch_size, sampler=train_sampler, shuffle=False) return train_loader, validation_loader
def get_data(): dataset = args.name path = '../data/geometric/MNIST' trainset = MNISTSuperpixels(path, train=True) testset = MNISTSuperpixels(path, train=False) lenTrain = len(trainset) lenTest = len(testset) trainLoader = DataLoader(trainset[:lenTrain // 50], batch_size=1, shuffle=False) testloader = DataLoader(testset[:lenTest // 50], batch_size=1, shuffle=False) return trainLoader, testloader
def load_pyg(name, dataset_dir): """ Load PyG dataset objects. (More PyG datasets will be supported) Args: name (string): dataset name dataset_dir (string): data directory Returns: PyG dataset object """ dataset_dir = '{}/{}'.format(dataset_dir, name) if name in ['Cora', 'CiteSeer', 'PubMed']: dataset = Planetoid(dataset_dir, name) elif name[:3] == 'TU_': # TU_IMDB doesn't have node features if name[3:] == 'IMDB': name = 'IMDB-MULTI' dataset = TUDataset(dataset_dir, name, transform=T.Constant()) else: dataset = TUDataset(dataset_dir, name[3:]) elif name == 'Karate': dataset = KarateClub() elif 'Coauthor' in name: if 'CS' in name: dataset = Coauthor(dataset_dir, name='CS') else: dataset = Coauthor(dataset_dir, name='Physics') elif 'Amazon' in name: if 'Computers' in name: dataset = Amazon(dataset_dir, name='Computers') else: dataset = Amazon(dataset_dir, name='Photo') elif name == 'MNIST': dataset = MNISTSuperpixels(dataset_dir) elif name == 'PPI': dataset = PPI(dataset_dir) elif name == 'QM7b': dataset = QM7b(dataset_dir) else: raise ValueError('{} not support'.format(name)) return dataset
def GCN(dataset, params, num_pre_epochs, num_epochs, MonteSize, PruningTimes, width, lr, savepath): Batch_size = int(params[0]) for Monte_iter in range(MonteSize): # Data best_loss = float('inf') # best test loss start_epoch = 0 # start from epoch 0 or last checkpoint epoch TrainConvergence = [] TestConvergence = [] # model root = '/git/data/GraphData/' + dataset if dataset == 'Cora': model_name = "PruningGCN" datasetroot = Planetoid(root=root, name=dataset).shuffle() trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) model_to_save = './checkpoint/{}-{}-param_{}_{}_{}_{}-Monte_{}-ckpt.pth'.format( dataset, model_name, params[0], params[1], params[2], params[3], Monte_iter) if Monte_iter == 0: if resume == True and os.path.exists(model_to_save): [ OptimizedNet, NewNetworksize, TrainConvergence, TestConvergence, start_epoch ] = ResumeModel(model_to_save) if start_epoch >= num_epochs - 1: continue else: net = Net(datasetroot, width) #net.apply(weight_reset) elif dataset == 'ENZYMES' or dataset == 'MUTAG': model_name = "topk_pool_Net" datasetroot = TUDataset(root, name=dataset) trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) model_to_save = './checkpoint/{}-{}-param_{}_{}_{}_{}-ckpt.pth'.format( dataset, model_name, params[0], params[1], params[2], params[3]) if Monte_iter == 0: if resume == True and os.path.exists(model_to_save): [ OptimizedNet, NewNetworksize, TrainConvergence, TestConvergence, start_epoch ] = ResumeModel(model_to_save) if start_epoch >= num_epochs - 1: continue else: net = topk_pool_Net(datasetroot, width) elif dataset == 'MNIST': datasetroot = MNISTSuperpixels(root=root, transform=T.Cartesian()).shuffle() trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) model_name = 'SPlineNet' model_to_save = './checkpoint/{}-{}-param_{}_{}_{}_{}-Mon_{}-ckpt.pth'.format( dataset, model_name, params[0], params[1], params[2], params[3], Monte_iter) if resume == "True" and os.path.exists(model_to_save): [net, TrainConvergence, TestConvergence, start_epoch] = ResumeModel(model_to_save) if start_epoch >= num_epochs - 1: continue else: #net=Net(datasetroot,width) net = SPlineNet(datasetroot, width) elif dataset == 'CIFAR10': if resume == "True" and os.path.exists(model_to_save): [net, TrainConvergence, TestConvergence, start_epoch] = ResumeModel(model_to_save) if start_epoch >= num_epochs - 1: continue else: net = getattr(CIFAR10_resnet, 'Resnet20_CIFAR10')(params[1]) else: raise Exception( "The dataset is:{}, it isn't existed.".format(dataset)) if Monte_iter == 0 and start_epoch == 0: print('Let\'s use', torch.cuda.device_count(), 'GPUs!') net = DataParallel(net) device = torch.device( 'cuda' if torch.cuda.is_available() else 'cpu') net = net.to(device) #cudnn.benchmark = True criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4) logging( 'Batch size: {}, Number of layers:{} ConCoeff: {}, CutoffCoffi:{}, MonteSize:{}' .format(params[0], params[1], params[2], params[3], Monte_iter)) for epoch in range(num_pre_epochs): PreTrainLoss = train(trainloader, net, optimizer, criterion) print('\nEpoch: {}, Average pre-tain loss: {:.4f} \n'.format( epoch, PreTrainLoss[0])) NewNetworksize = RetainNetworkSize(net, params[2]) del net #NewNetworksize=width for pruningIter in range(PruningTimes): if pruningIter > 0: [ OptimizedNet, NewNetworksize, TrainConvergence, TestConvergence, start_epoch ] = ResumeModel(model_to_save) elif dataset == 'Cora' and start_epoch == 0: OptimizedNet = Net(datasetroot, NewNetworksize[0:-1]) elif dataset == 'ENZYMES' and start_epoch == 0: NewNetworkSizeAdjust = NewNetworksize[0:-1] NewNetworkSizeAdjust[0] = width[0] - 1 OptimizedNet = topk_pool_Net(datasetroot, NewNetworkSizeAdjust) OptimizedNet = DataParallel(OptimizedNet) device = torch.device( 'cuda' if torch.cuda.is_available() else 'cpu') OptimizedNet = OptimizedNet.to(device) cudnn.benchmark = True criterionNew = nn.CrossEntropyLoss() optimizerNew = optim.SGD(OptimizedNet.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4) criterion = nn.CrossEntropyLoss() for epoch in range(start_epoch, num_epochs): TrainLoss = train(trainloader, OptimizedNet, optimizerNew, criterionNew) print('\n Epoch: {}, Average tain loss: {:.4f} \n'.format( epoch, TrainLoss[0])) TrainConvergence.append(statistics.mean(TrainLoss)) NewNetworksize = RetainNetworkSize(OptimizedNet, params[2]) TestConvergence.append( statistics.mean(test(testloader, OptimizedNet, criterion))) # save model if TestConvergence[epoch] < best_loss: logging('Saving..') state = { 'net': OptimizedNet.module, 'TrainConvergence': TrainConvergence, 'TestConvergence': TestConvergence, 'epoch': num_epochs, 'NewNetworksize': NewNetworksize[0:-1], } if not os.path.isdir('checkpoint'): os.mkdir('checkpoint') torch.save(state, model_to_save) best_loss = TestConvergence[epoch] ## save recurrence plots """if epoch%20==0: save_recurrencePlots_file="../Results/RecurrencePlots/RecurrencePlots_{}_{}_BatchSize{} \_ConCoeffi{}_epoch{}.png".format(dataset, model_name,params[0],params[1],epoch) save_recurrencePlots(net,save_recurrencePlots_file)""" FileName = "{}-{}-param_{}_{}_{}_{}-monte_{}".format( dataset, model_name, params[0], params[1], params[2], params[3], Monte_iter) np.save(savepath + 'TrainConvergence-' + FileName, TrainConvergence) #np.save(savepath+'TestConvergence-'+FileName,TestConvergence) #torch.cuda.empty_cache() print_nvidia_useage() if return_output == True: return TestConvergence[-1], net.module.fc.weight else: pass
import os.path as osp import torch import torch.nn.functional as F from torch_geometric.datasets import MNISTSuperpixels import torch_geometric.transforms as T from torch_geometric.data import DataLoader from torch_geometric.nn import SplineConv, voxel_grid, max_pool, max_pool_x path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST') transform = T.Cartesian(cat=False) train_dataset = MNISTSuperpixels(path, True, transform=transform) test_dataset = MNISTSuperpixels(path, False, transform=transform) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=64) d = train_dataset class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = SplineConv(d.num_features, 32, dim=2, kernel_size=5) self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5) self.conv3 = SplineConv(64, 64, dim=2, kernel_size=5) self.fc1 = torch.nn.Linear(4 * 64, 128) self.fc2 = torch.nn.Linear(128, d.num_classes) def forward(self, data): data.x = F.elu(self.conv1(data.x, data.edge_index, data.edge_attr)) cluster = voxel_grid(data.pos, data.batch, size=5, start=0, end=28) data.edge_attr = None
def prepare_data(self): path = osp.join( osp.dirname(osp.realpath(__file__)), "..", "..", "data", self.NAME ) self.train_dataset = MNISTSuperpixels(path, True, transform=self._transform) self.test_dataset = MNISTSuperpixels(path, False, transform=self._transform)
# rr, cc = draw.line(start_y, start_x, end_y, end_x) # image[rr, cc] = [1, 0, 0] # # Draw nodes # for i in range(position.size(0)): # x, y = position[i] # rr, cc = draw.circle(y, x, 10, shape=shape) # image[rr, cc] = [1, 0, 0] # rr, cc = draw.circle(y, x, 9, shape=shape) # image[rr, cc] = [input[i], input[i], input[i]] # return image image_dataset = MNIST('/tmp/MNIST', train=True, download=True) graph_dataset = MNISTSuperpixels('~/MNISTSuperpixels', train=True) scale = 32 rescale = 4 offset = torch.FloatTensor([1, 1]) example = 35 image, _ = image_dataset[example] image = np.array(image) data = graph_dataset[example] input, adj, position = data['input'], data['adj'], data['position'] print(data['target']) # adj = grid(torch.Size([28, 28]), connectivity=8) # position = grid_position(torch.Size([28, 28])) # input = image.flatten() / 255.0
def get_dataset(name, sparse=True, feat_str="deg+ak3+reall", root=None): if root is None or root == '': path = osp.join(osp.expanduser('~'), 'pyG_data', name) else: path = osp.join(root, name) degree = feat_str.find("deg") >= 0 onehot_maxdeg = re.findall("odeg(\d+)", feat_str) onehot_maxdeg = int(onehot_maxdeg[0]) if onehot_maxdeg else None k = re.findall("an{0,1}k(\d+)", feat_str) k = int(k[0]) if k else 0 groupd = re.findall("groupd(\d+)", feat_str) groupd = int(groupd[0]) if groupd else 0 remove_edges = re.findall("re(\w+)", feat_str) remove_edges = remove_edges[0] if remove_edges else 'none' centrality = feat_str.find("cent") >= 0 coord = feat_str.find("coord") >= 0 pre_transform = FeatureExpander(degree=degree, onehot_maxdeg=onehot_maxdeg, AK=k, centrality=centrality, remove_edges=remove_edges, group_degree=groupd).transform if 'MNIST' in name or 'CIFAR' in name: if name == 'MNIST_SUPERPIXEL': train_dataset = MNISTSuperpixels(path, True, pre_transform=pre_transform, transform=T.Cartesian()) test_dataset = MNISTSuperpixels(path, False, pre_transform=pre_transform, transform=T.Cartesian()) else: train_dataset = ImageDataset(path, name, True, pre_transform=pre_transform, coord=coord, processed_file_prefix="data_%s" % feat_str) test_dataset = ImageDataset(path, name, False, pre_transform=pre_transform, coord=coord, processed_file_prefix="data_%s" % feat_str) dataset = (train_dataset, test_dataset) elif 'QM9' in name: dataset = QM9Ext(path, pre_transform=pre_transform, processed_filename="data_%s.pt" % feat_str) elif 'ModelNet' in name: pre_transform = FeatureExpander( degree=degree, onehot_maxdeg=onehot_maxdeg, AK=k, centrality=centrality, remove_edges=remove_edges, group_degree=groupd).cloud_point_transform train_dataset = ModelNetExT(path, train=True, pre_transform=pre_transform, processed_file_prefix="data_%s" % feat_str) test_dataset = ModelNetExT(path, train=True, pre_transform=pre_transform, processed_file_prefix="data_%s" % feat_str) dataset = (train_dataset, test_dataset) elif 'TOSCA' in name: # pre_transform = FeatureExpander( # degree=degree, onehot_maxdeg=onehot_maxdeg, AK=k, # centrality=centrality, remove_edges=remove_edges, # group_degree=groupd).cloud_point_transform dataset = TOSCAEXT(path, pre_transform=pre_transform, processed_file_prefix="data_%s" % feat_str) else: dataset = TUDatasetExt(path, name, pre_transform=pre_transform, use_node_attr=True, processed_filename="data_%s.pt" % feat_str) dataset.data.edge_attr = None return dataset
def prepare_data(self): self.train_dataset = MNISTSuperpixels(args.data_fp, True, pre_transform=T.Compose([remove_self_loops, T.Polar()])) self.test_dataset = MNISTSuperpixels(args.data_fp, False, pre_transform=T.Compose([remove_self_loops, T.Polar()]))
# print(data.num_nodes) # print(data.num_edges) # print(data.num_features) # print(data.is_directed()) # Loading benchmark datasets from torch_geometric.data import InMemoryDataset from torch_geometric.datasets import MNISTSuperpixels, TUDataset, ModelNet dataset = 'Mnist' path = '../data/geometric/MNIST' dataset = MNISTSuperpixels(path, train=True) split = int(len(dataset) * .9) # Randomly permute dataset before splitting dataset = dataset.shuffle() print(dataset.num_classes) print(dataset.num_features) print(dataset[0]) train = dataset[:split] test = dataset[split:] # dataset = ModelNet(root='../data/ModelNet', name='10') # print(len(dataset)) # print(dataset.num_classes)
import os.path as osp import torch import torch.nn.functional as F from torch_geometric.datasets import MNISTSuperpixels from torch_geometric.data import DataListLoader from torch_geometric.nn import GCNConv, global_mean_pool, DataParallel path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST') dataset = MNISTSuperpixels(path).shuffle() loader = DataListLoader(dataset, batch_size=512, shuffle=True) class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = GCNConv(dataset.num_features, 64) self.lin = torch.nn.Linear(64, dataset.num_classes) def forward(self, data): print('Inside Model: num graphs: {}, device: {}'.format( data.num_graphs, data.batch.device)) x = self.conv1(data.x, data.edge_index) x = global_mean_pool(x, data.batch) return F.log_softmax(self.lin(x), dim=-1) model = Net() if torch.cuda.device_count() > 1: print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
def main(args, kwargs): dataset = args.data.upper() if dataset in ['CORA', 'CITESEER', 'PUBMED']: path = '../data/geometric/' + dataset print('Using {} dataset'.format(dataset)) dataset = Planetoid(path, dataset, T.NormalizeFeatures()) elif dataset == 'MNIST': path = '../data/geometric/' + dataset print('Using {} dataset'.format(dataset)) dataset = MNISTSuperpixels(path, dataset, T.NormalizeFeatures()) else: sys.exit( "You must choose one of the 'Planetoid' datasets (CORA, CITESEER, or PUBMED)." ) data = dataset[0] # Store the original adjacnecy matrix (for later calculation of edge prediction accuracy) adj_original = torch.tensor(np.asarray(get_adjacency(data).toarray(), dtype=np.float32), requires_grad=False) channels = 16 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = kwargs[args.model](Encoder(dataset.num_features, channels)).to(device) data.train_mask = data.val_mask = data.test_mask = data.y = None data = model.split_edges(data) x, edge_index = data.x.to(device), data.edge_index.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.01) # Create defaultdict of lists to be used for keeping running tally of accuracy, AUC, and AP scores results = defaultdict(list) def train(): model.train() optimizer.zero_grad() # Produces N * channels vector z = model.encode(x, edge_index) ''' TESTING DECODE_INDICES ''' # pos_decoded = model.decode_indices(z, data.train_pos_edge_index) # print('z shape: ', z.detach().numpy().shape) # print('pos edge index shape: ', data.train_pos_edge_index.detach().numpy().shape) # print('decode_indices_pos: ', pos_decoded.detach().numpy()) # print('decode_indices_pos shape: ', pos_decoded.detach().numpy().shape) # neg_edge_index = negative_sampling(data.train_pos_edge_index, z.size(0)) # neg_decoded = model.decode_indices(z, neg_edge_index) # print('decode_indices_neg: ', neg_decoded.detach().numpy()) # print('decode_indices_neg shape: ', neg_decoded.detach().numpy().shape) ''' END TESTING ''' ''' Testing Kernel Similarity ''' # g1 = nx.from_scipy_sparse_matrix(adj_original) # print(g1) # from sklearn.metrics.pairwise import rbf_kernel # adj_reconstructed = model.decode(z, sigmoid=True).detach().numpy() # adj_reconstructed = (adj_reconstructed > .5).astype(int) # similarity = kernel_similarity(adj_original, adj_reconstructed, rbf_kernel) # print(similarity) ''' END TESTING ''' if args.loss == 'bce': loss = model.recon_loss(z, data.train_pos_edge_index) results['recon_loss'].append(loss) kl_loss = model.kl_loss() results['kl'].append(kl_loss) loss = loss + 0.001 * kl_loss elif args.loss == 'newbce': numNodes = len(data['x']) loss = model.new_recon_loss(z, edge_index, num_nodes=numNodes, num_channels=channels, adj_original=adj_original) results['recon_loss'].append(loss) kl_loss = model.kl_loss() results['kl'].append(kl_loss) loss = loss + 0.001 * kl_loss elif args.loss == 'l2': loss = model.recon_loss_l2(z, adj_original) results['recon_loss'].append(loss) kl_loss = model.kl_loss() results['kl'].append(kl_loss) loss = loss + 0.001 * kl_loss elif args.loss == 'anneal': # anneal_weight = (1.0 - args.kl_weight) / (len(data.train_pos_edge_index) ) # linear annealing anneal_rate = epoch / args.kl_warmup if epoch < args.kl_warmup else 1. kl_weight = min(1.0, anneal_rate) # Starting annealing # kl_weight = args.kl_start # anneal_rate = (1.0 - args.kl_start) / (args.kl_warmup * (len(data))) # kl_weight = min(1.0, kl_weight + anneal_rate) print(kl_weight) loss = model.recon_loss(z, data.train_pos_edge_index) results['recon_loss'].append(loss) kl_loss = model.kl_loss() * kl_weight results['kl'].append(kl_loss) loss = loss + 0.001 * (kl_loss) results['loss'].append(loss) loss.backward() optimizer.step() def validate(pos_edge_index, neg_edge_index): model.eval() with torch.no_grad(): z = model.encode(x, edge_index) # accuracy = model.get_accuracy(z, pos_edge_index, neg_edge_index, adj_original) # accuracy = model.get_accuracy_new(z, adj_original) # print(accuracy) auc, ap = model.test(z, pos_edge_index, neg_edge_index) # return accuracy, auc, ap return auc, ap def test(pos_edge_index, neg_edge_index): model.eval() with torch.no_grad(): z = model.encode(x, edge_index) # accuracy = model.get_accuracy(z, pos_edge_index, neg_edge_index, adj_original) # accuracy = model.get_accuracy_new(z, adj_original) auc, ap = model.test(z, pos_edge_index, neg_edge_index) # return accuracy, auc, ap return auc, ap for epoch in range(1, args.epochs): train() # Run on validation edges # accuracy, auc, ap = validate(data.val_pos_edge_index, data.val_neg_edge_index) auc, ap = validate(data.val_pos_edge_index, data.val_neg_edge_index) # results['acc_val'].append(accuracy) results['auc_val'].append(auc) results['ap_val'].append(ap) # Print AUC and AP on validation data epochs if epoch % 10 == 0: # print('Val Epoch : {:03d}, ACC: {:.4f}, AUC: {:.4f}, AP: {:.4f}'.format(epoch, accuracy, auc, ap)) print('Val Epoch : {:03d}, AUC: {:.4f}, AP: {:.4f}'.format( epoch, auc, ap)) # Evalulate on heldout test edges for every epoch which is a multiple of the test_freq argument if epoch % args.test_freq == 0: # accuracy, auc, ap = test(data.test_pos_edge_index, data.test_neg_edge_index) auc, ap = test(data.test_pos_edge_index, data.test_neg_edge_index) # results['acc_test'].append(accuracy) results['auc_test'].append(auc) results['ap_test'].append(ap) # print('Test Epoch: {:03d}, ACC: {:.4f}, AUC: {:.4f}, AP: {:.4f}'.format(epoch, accuracy, auc, ap)) print('Test Epoch: {:03d}, AUC: {:.4f}, AP: {:.4f}'.format( epoch, auc, ap)) # Evaluate on held-out test edges # auc, ap = test(data.test_pos_edge_index, data.test_neg_edge_index) # print('Test AUC: {:.4f}, Test AP: {:.4f}'.format(auc, ap)) # Pickle results if args.save: if args.notes is None: modelPath = '../models/' + args.data + '_RESULTS.p' plotPath = '../figures/geometric/' + args.data + '_RESULTS.png' lossPath = '../figures/geometric/' + args.data + '_LOSSES.png' else: modelPath = '../models/' + args.notes + '_' + args.data + '_RESULTS.p' plotPath = '../figures/geometric/' + args.notes + '_' + args.data + '_RESULTS.png' lossPath = '../figures/geometric/' + args.notes + '_' + args.data + '_LOSSES.png' pkl.dump(results, open(modelPath, 'wb')) plot_results(pkl.load(open(modelPath, 'rb')), path=plotPath, loss=args.loss, anneal=args.kl_warmup) plot_losses(pkl.load(open(modelPath, 'rb')), path=lossPath, loss=args.loss, anneal=args.kl_warmup)
def load_dataset(args): # automatic data loading and splitting transform = add_zeros if args.dataset == 'ogbg-ppa' else None cls_criterion = get_loss_function(args.dataset) idx2word_mapper = None if args.dataset == 'mnist': train_data = MNISTSuperpixels(root='dataset', train=True, transform=T.Polar()) dataset = train_data dataset.name = 'mnist' dataset.eval_metric = 'acc' validation_data = [] test_data = MNISTSuperpixels(root='dataset', train=False, transform=T.Polar()) train_data = list(train_data) test_data = list(test_data) elif args.dataset == 'QM9': # Contains 19 targets. Use only the first 12 (0-11) QM9_VALIDATION_START = 110000 QM9_VALIDATION_END = 120000 dataset = QM9(root='dataset', transform=ExtractTargetTransform(args.target)).shuffle() dataset.name = 'QM9' dataset.eval_metric = 'mae' train_data = dataset[:QM9_VALIDATION_START] validation_data = dataset[QM9_VALIDATION_START:QM9_VALIDATION_END] test_data = dataset[QM9_VALIDATION_END:] train_data = list(train_data) validation_data = list(validation_data) test_data = list(test_data) elif args.dataset == 'zinc': train_data = ZINC(root='dataset', subset=True, split='train') dataset = train_data dataset.name = 'zinc' validation_data = ZINC(root='dataset', subset=True, split='val') test_data = ZINC(root='dataset', subset=True, split='test') dataset.eval_metric = 'mae' train_data = list(train_data) validation_data = list(validation_data) test_data = list(test_data) elif args.dataset in [ 'ogbg-molhiv', 'ogbg-molpcba', 'ogbg-ppa', 'ogbg-code2' ]: dataset = PygGraphPropPredDataset(name=args.dataset, transform=transform) if args.dataset == 'obgb-code2': seq_len_list = np.array([len(seq) for seq in dataset.data.y]) max_seq_len = args.max_seq_len num_less_or_equal_to_max = np.sum( seq_len_list <= args.max_seq_len) / len(seq_len_list) print( f'Target sequence less or equal to {max_seq_len} is {num_less_or_equal_to_max}%.' ) split_idx = dataset.get_idx_split() # The following is only used in the evaluation of the ogbg-code classifier. if args.dataset == 'ogbg-code2': vocab2idx, idx2vocab = get_vocab_mapping( [dataset.data.y[i] for i in split_idx['train']], args.num_vocab) # specific transformations for the ogbg-code dataset dataset.transform = transforms.Compose([ augment_edge, lambda data: encode_y_to_arr(data, vocab2idx, args.max_seq_len) ]) idx2word_mapper = partial(decode_arr_to_seq, idx2vocab=idx2vocab) train_data = list(dataset[split_idx["train"]]) validation_data = list(dataset[split_idx["valid"]]) test_data = list(dataset[split_idx["test"]]) return dataset, train_data, validation_data, test_data, cls_criterion, idx2word_mapper
def TrainingNet(dataset,modelName,params,num_pre_epochs,num_epochs,NumCutoff,optimizerName,LinkPredictionMethod,MonteSize,savepath): Batch_size=params[0] VectorPairs=params[4] StartTopoCoeffi=params[5] WeightCorrectionCoeffi=params[6] interval=params[7] root='/git/data/GraphData/'+dataset TestAccs=[] for Monte_iter in range(MonteSize): # Data NewNetworkSizeAdjust=[] WeightsDynamicsEvolution=[] trainValRatio=[0.2,0.4] # model if dataset=='Cora' or dataset =='Citeseer' or dataset =='Pubmed': datasetroot= Planetoid(root=root, name=dataset,transform =T.NormalizeFeatures()).shuffle() trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) """ train_mask, val_mask,test_mask=DataSampler(trainValRatio,datasetroot.data.num_nodes) DataMask={} DataMask['train_mask']=train_mask DataMask['val_mask']=val_mask DataMask['test_mask']=test_mask trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)""" num_features=datasetroot.num_features num_classes=datasetroot.num_classes criterion = nn.CrossEntropyLoss() elif dataset =="CoraFull": datasetroot = CoraFull(root=root,transform =T.NormalizeFeatures()).shuffle() """train_mask, val_mask,test_mask=DataSampler(trainValRatio,datasetroot.data.num_nodes) DataMask={} DataMask['train_mask']=train_mask DataMask['val_mask']=val_mask DataMask['test_mask']=test_mask""" criterion = nn.CrossEntropyLoss() num_features=datasetroot.num_features num_classes=datasetroot.num_classes trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=False) elif dataset=='ENZYMES' or dataset=='MUTAG': datasetroot=TUDataset(root,name=dataset,use_node_attr=True) trainloader = DataLoader(datasetroot, batch_size=Batch_size, shuffle=True) num_features=datasetroot.num_features num_classes=datasetroot.num_classes elif dataset =="PPI": train_dataset = PPI(root, split='train') val_dataset = PPI(root, split='val') test_dataset = PPI(root, split='test') trainloader = DataListLoader(train_dataset, batch_size=Batch_size, shuffle=True) valloader = DataListLoader(val_dataset, batch_size=100, shuffle=False) testloader = DataListLoader(test_dataset, batch_size=100, shuffle=False) num_classes=train_dataset.num_classes num_features=train_dataset.num_features criterion = torch.nn.BCEWithLogitsLoss() elif dataset =="Reddit": datasetroot=Reddit(root) trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=2, shuffle=False) elif dataset=="Amazon": datasetroot=Amazon(root, "Photo", transform=None, pre_transform=None) trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) elif dataset=='MNIST': datasetroot = MNISTSuperpixels(root=root, transform=T.Cartesian()) trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) elif dataset=='CIFAR10': pass else: raise Exception("Input wrong datatset!!") width=ContractionLayerCoefficients(num_features,*params[1:3]) net =ChooseModel(modelName,num_features,num_classes,width) FileName="{}-{}-param_{}_{}_{}_{}-monte_{}".format(dataset,modelName,interval,WeightCorrectionCoeffi,StartTopoCoeffi,VectorPairs,Monte_iter) print('Let\'s use', torch.cuda.device_count(), 'GPUs!') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') criterion = criterion.to(device) net = DataParallel(net) net = net.to(device) optimizer = getattr(optim,optimizerName)(net.parameters(), lr=params[3], momentum=0.9, weight_decay=5e-4) model_to_save='./checkpoint/{}-{}-param_{}_{}_{}_{}-ckpt.pth'.format(dataset,modelName,params[0],params[1],params[5],params[4]) if resume=="True" and os.path.exists(model_to_save): [net,optimizer,TrainConvergence,TestConvergence,Acc]=ResumeModel(net,optimizer,model_to_save) start_epoch=len(TrainConvergence) else: start_epoch = 0 # start from epoch 0 or last checkpoint epoch #cudnn.benchmark = True logging('dataset:{}, Batch size: {}, Number of layers:{} ConCoeff: {}, LR:{}, MonteSize:{}'.format(dataset, params[0], params[1],params[2],params[3],Monte_iter)) mark="{}{}Convergence/DiagElement-{}".format(savepath,dataset,FileName) markweights="{}{}Convergence/WeightChanges-{}".format(savepath,dataset,FileName) PreTrainConvergence,PreTestConvergence,PreAcc=TrainPart(start_epoch,num_pre_epochs,num_classes, trainloader,net,optimizer,criterion,NumCutoff,LinkPredictionMethod,VectorPairs,WeightCorrectionCoeffi,StartTopoCoeffi,mark,markweights,model_to_save,False) print('dataset: {}, model name:{}, epoch:{},Pre-train error:{}; Pre-test error:{}; test acc:{}'.format(dataset,modelName,num_pre_epochs,PreTrainConvergence[-1],PreTestConvergence[-1],PreAcc)) NewNetworksize=RetainNetworkSize(net,params[2]) OptimizedNet=ChooseModel(modelName,num_features,num_classes,NewNetworksize[0:-1]) NewNetworksize.insert(0,num_features) NewNetworkSizeAdjust.append(NewNetworksize[0:-1]) print(NewNetworkSizeAdjust) #OptimizedNet.apply(init_weights) OptimizedNet = DataParallel(OptimizedNet) OptimizedNet = OptimizedNet.to(device) cudnn.benchmark = True # Begin Pre training if optimizerName =="SGD": optimizerNew = getattr(optim,optimizerName)(OptimizedNet.parameters(), lr=params[3], momentum=0.9, weight_decay=5e-4) elif optimizerName =="Adam": optimizerNew = getattr(optim,optimizerName)(OptimizedNet.parameters(), lr=params[3], betas=(0.9, 0.999), eps=1e-08, weight_decay=5e-4, amsgrad=False) TrainConvergence,TestConvergence,TestAcc=TrainPart(start_epoch,num_epochs,datasetroot.num_classes, trainloader,OptimizedNet,optimizerNew,criterion, NumCutoff,LinkPredictionMethod,VectorPairs,WeightCorrectionCoeffi,StartTopoCoeffi,mark,markweights,model_to_save,True) np.save("{}/{}Convergence/AlgebraicConectivityTrainConvergence-{}".format(savepath,dataset,FileName),TrainConvergence) np.save("{}/{}Convergence/AlgebraicConectivityTestConvergence-{}".format(savepath,dataset,FileName),TestConvergence) #np.save("{}/{}Convergence/NewNetworkSizeAdjust-{}".format(savepath,dataset,FileName),NewNetworkSizeAdjust) #torch.cuda.empty_cache() print('dataset: {}, model name:{}, resized network size: {}, the train error: {},test error: {}, test acc:{}\n'.format(dataset,modelName,NewNetworksize[0:-1],num_epochs,TrainConvergence[-1],TestConvergence[-1],TestAcc)) np.save("{}/{}Convergence/AlgebraicConectivityMeanTestAccs-{}".format(savepath,dataset,FileName),TestAccs.append(TestAcc)) TestAccs.append(TestAcc) print_nvidia_useage()
def GCN(dataset, params, Epochs, MonteSize, width, lr, savepath): Batch_size = int(params[0]) for Monte_iter in range(MonteSize): # Data best_loss = float('inf') # best test loss start_epoch = 0 # start from epoch 0 or last checkpoint epoch TrainConvergence = [] TestConvergence = [] # model root = '/data/GraphData/' + dataset if dataset == 'Cora': model_name = "GCN3" datasetroot = Planetoid(root=root, name=dataset).shuffle() trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) model_to_save = './checkpoint/{}-{}-param_{}_{}-Mon_{}-ckpt.pth'.format( dataset, model_name, params[0], params[1], Monte_iter) if resume and os.path.exists(model_to_save): [net, TrainConvergence, TestConvergence, start_epoch] = ResumeModel(model_to_save) if start_epoch >= Epochs - 1: continue else: net = Net(datasetroot, width) elif dataset == 'ENZYMES' or dataset == 'MUTAG': model_name = "topk_pool_Net" root = '/data/GraphData' + dataset datasetroot = TUDataset(root, name=dataset) trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) model_to_save = './checkpoint/{}-{}-param_{}_{}-Mon_{}-ckpt.pth'.format( dataset, model_name, params[0], params[1], Monte_iter) if resume and os.path.exists(model_to_save): [net, TrainConvergence, TestConvergence, start_epoch] = ResumeModel(model_to_save) if start_epoch >= Epochs - 1: continue else: net = topk_pool_Net(datasetroot, width) elif dataset == 'MNIST': datasetroot = MNISTSuperpixels(root='/data/GraphData/' + dataset, transform=T.Cartesian()).shuffle() trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) model_name = 'SPlineNet' model_to_save = './checkpoint/{}-{}-param_{}_{}-Mon_{}-ckpt.pth'.format( dataset, model_name, params[0], params[1], Monte_iter) if resume and os.path.exists(model_to_save): [net, TrainConvergence, TestConvergence, start_epoch] = ResumeModel(model_to_save) if start_epoch >= Epochs - 1: continue else: #net=Net(datasetroot,width) net = SPlineNet(datasetroot, width) elif dataset == 'CIFAR10': if resume and os.path.exists(model_to_save): [net, TrainConvergence, TestConvergence, start_epoch] = ResumeModel(model_to_save) if start_epoch >= Epochs - 1: continue else: net = getattr(CIFAR10_resnet, 'Resnet20_CIFAR10')(params[1]) else: raise Exception( "The dataset is:{}, it isn't existed.".format(dataset)) print('Let\'s use', torch.cuda.device_count(), 'GPUs!') torch.cuda.is_available() net = DataParallel(net) device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') net = net.to(device) #cudnn.benchmark = True criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4) for epoch in range(start_epoch, start_epoch + Epochs): if epoch < Epochs: logging( 'Batch size: {},ConCoeff: {},MonteSize:{},epoch:{}'.format( params[0], params[1], Monte_iter, epoch)) TrainLoss = train(trainloader, net, optimizer, criterion) TrainConvergence.append(statistics.mean(TrainLoss)) TestConvergence.append( statistics.mean(test(testloader, net, criterion))) else: break if TestConvergence[epoch] < best_loss: logging('Saving..') state = { 'net': net.module, 'TrainConvergence': TrainConvergence, 'TestConvergence': TestConvergence, 'epoch': epoch, } if not os.path.isdir('checkpoint'): os.mkdir('checkpoint') torch.save(state, model_to_save) best_loss = TestConvergence[epoch] if not os.path.exists('./%s' % model_name): os.makedirs('./%s' % model_name) torch.save( net.module.state_dict(), './%s/%s_%s_%s_%s_%s_pretrain.pth' % (model_name, dataset, model_name, params[0], params[1], Epochs)) else: pass ## save recurrence plots if epoch % 20 == 0: save_recurrencePlots_file = "../Results/RecurrencePlots/RecurrencePlots_{}_{}_BatchSize{}_ConCoeffi{}_epoch{}.png".format( dataset, model_name, params[0], params[1], epoch) save_recurrencePlots(net, save_recurrencePlots_file) FileName = "{}-{}-param_{}_{}-monte_{}".format(dataset, model_name, params[0], params[1], Monte_iter) np.save(savepath + 'TrainConvergence-' + FileName, TrainConvergence) np.save(savepath + 'TestConvergence-' + FileName, TestConvergence) torch.cuda.empty_cache() print_nvidia_useage() if return_output == True: return TestConvergence[-1], net.module.fc.weight else: pass
def main(args): args = init(args) def pf(data): return data.y == args.num pre_filter = pf if args.num != -1 else None print("loading data") if (args.sparse_mnist): X = MNISTGraphDataset(args.dataset_path, args.num_hits, train=args.train, num=args.num) X_loaded = DataLoader(X, shuffle=True, batch_size=args.batch_size, pin_memory=True) else: if (args.gcnn): X = MNISTSuperpixels(args.dir_path, train=args.train, pre_transform=T.Cartesian(), pre_filter=pre_filter) X_loaded = tgDataLoader(X, shuffle=True, batch_size=args.batch_size) else: X = SuperpixelsDataset(args.dataset_path, args.num_hits, train=args.train, num=args.num) X_loaded = DataLoader(X, shuffle=True, batch_size=args.batch_size, pin_memory=True) print("loaded data") # model if (args.load_model): G = torch.load(args.model_path + args.name + "/G_" + str(args.start_epoch) + ".pt", map_location=args.device) D = torch.load(args.model_path + args.name + "/D_" + str(args.start_epoch) + ".pt", map_location=args.device) else: # G = Graph_Generator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.fn_hidden_size, args.fn_num_layers, args.mp_iters_gen, args.num_hits, args.gen_dropout, args.leaky_relu_alpha, hidden_node_size=args.hidden_node_size, int_diffs=args.int_diffs, pos_diffs=args.pos_diffs, gru=args.gru, batch_norm=args.batch_norm, device=device).to(args.device) if (args.gcnn): G = GaussianGenerator(args=deepcopy(args)).to(args.device) D = MoNet(args=deepcopy(args)).to(args.device) # D = Gaussian_Discriminator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.mp_hidden_size, args.mp_num_layers, args.num_iters, args.num_hits, args.dropout, args.leaky_relu_alpha, kernel_size=args.kernel_size, hidden_node_size=args.hidden_node_size, int_diffs=args.int_diffs, gru=GRU, batch_norm=args.batch_norm, device=device).to(args.device) else: # D = Graph_Discriminator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.fn_hidden_size, args.fn_num_layers, args.mp_iters_disc, args.num_hits, args.disc_dropout, args.leaky_relu_alpha, hidden_node_size=args.hidden_node_size, wgan=args.wgan, int_diffs=args.int_diffs, pos_diffs=args.pos_diffs, gru=args.gru, batch_norm=args.batch_norm, device=device).to(args.device) print("Generator") G = Graph_GAN(gen=True, args=deepcopy(args)).to(args.device) print("Discriminator") D = Graph_GAN(gen=False, args=deepcopy(args)).to(args.device) print("Models loaded") # optimizer if args.spectral_norm_gen: G_params = filter(lambda p: p.requires_grad, G.parameters()) else: G_params = G.parameters() if args.spectral_norm_gen: D_params = filter(lambda p: p.requires_grad, D.parameters()) else: D_params = D.parameters() if (args.optimizer == 'rmsprop'): G_optimizer = optim.RMSprop(G_params, lr=args.lr_gen) D_optimizer = optim.RMSprop(D_params, lr=args.lr_disc) elif (args.optimizer == 'adadelta'): G_optimizer = optim.Adadelta(G_params, lr=args.lr_gen) D_optimizer = optim.Adadelta(D_params, lr=args.lr_disc) elif (args.optimizer == 'acgd'): optimizer = ACGD(max_params=G_params, min_params=D_params, lr_max=args.lr_gen, lr_min=args.lr_disc, device=args.device) elif (args.optimizer == 'adam' or args.optimizer == 'None'): G_optimizer = optim.Adam(G_params, lr=args.lr_gen, weight_decay=5e-4, betas=(args.beta1, args.beta2)) D_optimizer = optim.Adam(D_params, lr=args.lr_disc, weight_decay=5e-4, betas=(args.beta1, args.beta2)) if (args.load_model): try: if (not args.optimizer == 'acgd'): G_optimizer.load_state_dict( torch.load(args.model_path + args.name + "/G_optim_" + str(args.start_epoch) + ".pt", map_location=args.device)) D_optimizer.load_state_dict( torch.load(args.model_path + args.name + "/D_optim_" + str(args.start_epoch) + ".pt", map_location=args.device)) else: optimizer.load_state_dict( torch.load(args.model_path + args.name + "/optim_" + str(args.start_epoch) + ".pt", map_location=args.device)) except: print("Error loading optimizer") print("optimizers loaded") if args.fid: C, mu2, sigma2 = evaluation.load(args, X_loaded) normal_dist = Normal( torch.tensor(0.).to(args.device), torch.tensor(args.sd).to(args.device)) lns = args.latent_node_size if args.latent_node_size else args.hidden_node_size args.noise_file_name = "num_samples_" + str( args.num_samples) + "_num_nodes_" + str( args.num_hits) + "_latent_node_size_" + str(lns) + "_sd_" + str( args.sd) + ".pt" if args.gcnn: args.noise_file_name = "gcnn_" + args.noise_file_name noise_file_names = listdir(args.noise_path) if args.noise_file_name not in noise_file_names: if (args.gcnn): torch.save( normal_dist.sample( (args.num_samples * 5, 2 + args.channels[0])), args.noise_path + args.noise_file_name) else: torch.save( normal_dist.sample((args.num_samples, args.num_hits, lns)), args.noise_path + args.noise_file_name) losses = {} if (args.load_model): try: losses['D'] = np.loadtxt(args.losses_path + args.name + "/" + "D.txt").tolist()[:args.start_epoch] losses['Dr'] = np.loadtxt(args.losses_path + args.name + "/" + "Dr.txt").tolist()[:args.start_epoch] losses['Df'] = np.loadtxt(args.losses_path + args.name + "/" + "Df.txt").tolist()[:args.start_epoch] losses['G'] = np.loadtxt(args.losses_path + args.name + "/" + "G.txt").tolist()[:args.start_epoch] if args.fid: losses['fid'] = np.loadtxt( args.losses_path + args.name + "/" + "fid.txt").tolist()[:args.start_epoch] if (args.gp): losses['gp'] = np.loadtxt(args.losses_path + args.name + "/" + "gp.txt").tolist()[:args.start_epoch] except: print("couldn't load losses") losses['D'] = [] losses['Dr'] = [] losses['Df'] = [] losses['G'] = [] if args.fid: losses['fid'] = [] if (args.gp): losses['gp'] = [] else: losses['D'] = [] losses['Dr'] = [] losses['Df'] = [] losses['G'] = [] if args.fid: losses['fid'] = [] if (args.gp): losses['gp'] = [] Y_real = torch.ones(args.batch_size, 1).to(args.device) Y_fake = torch.zeros(args.batch_size, 1).to(args.device) def train_D(data, gen_data=None, unrolled=False): if args.debug: print("dtrain") D.train() D_optimizer.zero_grad() run_batch_size = data.shape[0] if not args.gcnn else data.y.shape[0] if gen_data is None: gen_data = utils.gen(args, G, normal_dist, run_batch_size) if (args.gcnn): gen_data = utils.convert_to_batch(args, gen_data, run_batch_size) if args.augment: p = args.aug_prob if not args.adaptive_prob else losses['p'][-1] data = augment.augment(args, data, p) gen_data = augment.augment(args, gen_data, p) D_real_output = D(data.clone()) D_fake_output = D(gen_data) D_loss, D_loss_items = utils.calc_D_loss(args, D, data, gen_data, D_real_output, D_fake_output, run_batch_size, Y_real, Y_fake) D_loss.backward(create_graph=unrolled) D_optimizer.step() return D_loss_items def train_G(data): if args.debug: print("gtrain") G.train() G_optimizer.zero_grad() gen_data = utils.gen(args, G, normal_dist, args.batch_size) if (args.gcnn): gen_data = utils.convert_to_batch(args, gen_data, args.batch_size) if args.augment: p = args.aug_prob if not args.adaptive_prob else losses['p'][-1] gen_data = augment.augment(args, gen_data, p) if (args.unrolled_steps > 0): D_backup = deepcopy(D) for i in range(args.unrolled_steps - 1): train_D(data, gen_data=gen_data, unrolled=True) D_fake_output = D(gen_data) G_loss = utils.calc_G_loss(args, D_fake_output, Y_real) G_loss.backward() G_optimizer.step() if (args.unrolled_steps > 0): D.load(D_backup) return G_loss.item() def train_acgd(data): if args.debug: print("acgd train") D.train() G.train() optimizer.zero_grad() run_batch_size = data.shape[0] if not args.gcnn else data.y.shape[0] gen_data = utils.gen(args, G, normal_dist, run_batch_size) if (args.gcnn): gen_data = utils.convert_to_batch(args, gen_data, run_batch_size) if args.augment: p = args.aug_prob if not args.adaptive_prob else losses['p'][-1] data = utils.rand_translate(args, data, p) gen_data = utils.rand_translate(args, gen_data, p) D_real_output = D(data.clone()) D_fake_output = D(gen_data) D_loss, D_loss_items = utils.calc_D_loss(args, D, data, gen_data, D_real_output, D_fake_output, run_batch_size) optimizer.step(loss=D_loss) G.eval() with torch.no_grad(): G_loss = utils.calc_G_loss(args, D_fake_output) return D_loss_items, G_loss.item() def train(): k = 0 temp_ng = args.num_gen if (args.fid): losses['fid'].append( evaluation.get_fid(args, C, G, normal_dist, mu2, sigma2)) if (args.save_zero): save_outputs.save_sample_outputs(args, D, G, normal_dist, args.name, 0, losses) for i in range(args.start_epoch, args.num_epochs): print("Epoch %d %s" % ((i + 1), args.name)) Dr_loss = 0 Df_loss = 0 G_loss = 0 D_loss = 0 gp_loss = 0 lenX = len(X_loaded) for batch_ndx, data in tqdm(enumerate(X_loaded), total=lenX): data = data.to(args.device) if (args.gcnn): data.pos = (data.pos - 14) / 28 row, col = data.edge_index data.edge_attr = (data.pos[col] - data.pos[row]) / (2 * args.cutoff) + 0.5 if (not args.optimizer == 'acgd'): if (args.num_critic > 1): D_loss_items = train_D(data) D_loss += D_loss_items['D'] Dr_loss += D_loss_items['Dr'] Df_loss += D_loss_items['Df'] if (args.gp): gp_loss += D_loss_items['gp'] if ((batch_ndx - 1) % args.num_critic == 0): G_loss += train_G(data) else: if (batch_ndx == 0 or (batch_ndx - 1) % args.num_gen == 0): D_loss_items = train_D(data) D_loss += D_loss_items['D'] Dr_loss += D_loss_items['Dr'] Df_loss += D_loss_items['Df'] if (args.gp): gp_loss += D_loss_items['gp'] G_loss += train_G(data) else: D_loss_items, G_loss_item = train_acgd(data) D_loss += D_loss_items['D'] Dr_loss += D_loss_items['Dr'] Df_loss += D_loss_items['Df'] G_loss += G_loss_item if args.bottleneck: if (batch_ndx == 10): return losses['D'].append(D_loss / (lenX / args.num_gen)) losses['Dr'].append(Dr_loss / (lenX / args.num_gen)) losses['Df'].append(Df_loss / (lenX / args.num_gen)) losses['G'].append(G_loss / (lenX / args.num_critic)) if (args.gp): losses['gp'].append(gp_loss / (lenX / args.num_gen)) print("d loss: " + str(losses['D'][-1])) print("g loss: " + str(losses['G'][-1])) print("dr loss: " + str(losses['Dr'][-1])) print("df loss: " + str(losses['Df'][-1])) if (args.gp): print("gp loss: " + str(losses['gp'][-1])) gloss = losses['G'][-1] drloss = losses['Dr'][-1] dfloss = losses['Df'][-1] dloss = (drloss + dfloss) / 2 if (args.bgm): if (i > 20 and gloss > dloss + args.bag): print("num gen upping to 10") args.num_gen = 10 else: print("num gen normal") args.num_gen = temp_ng elif (args.gom): if (i > 20 and gloss > dloss + args.bag): print("G loss too high - training G only") j = 0 print("starting g loss: " + str(gloss)) print("starting d loss: " + str(dloss)) while (gloss > dloss + args.bag * 0.5): print(j) gloss = 0 for l in tqdm(range(lenX)): gloss += train_G() gloss /= lenX print("g loss: " + str(gloss)) print("d loss: " + str(dloss)) losses['D'].append(dloss * 2) losses['Dr'].append(drloss) losses['Df'].append(dfloss) losses['G'].append(gloss) if (j % 5 == 0): save_outputs.save_sample_outputs(args, D, G, normal_dist, args.name, i + 1, losses, k=k, j=j) j += 1 k += 1 elif (args.rd): if (i > 20 and gloss > dloss + args.bag): print("gloss too high, resetting D params") D.reset_params() if ((i + 1) % 5 == 0): optimizers = optimizer if args.optimizer == 'acgd' else ( D_optimizer, G_optimizer) save_outputs.save_models(args, D, G, optimizers, args.name, i + 1) if (args.fid and (i + 1) % 1 == 0): losses['fid'].append( evaluation.get_fid(args, C, G, normal_dist, mu2, sigma2)) if ((i + 1) % 5 == 0): save_outputs.save_sample_outputs(args, D, G, normal_dist, args.name, i + 1, losses) train()
def main(): dataset = MNISTSuperpixels(root='~/MNISTSuperpixels', train=True) dataset_test = MNISTSuperpixels(root='~/MNISTSuperpixels', train=False) class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = GCNConv(dataset.num_node_features, 32) self.conv2 = GCNConv(32, 16) self.conv3 = GCNConv(16, dataset.num_classes) #self.conv1 = SAGEConv(dataset.num_node_features, 32) #self.conv2 = SAGEConv(32, 16) #self.conv3 = SAGEConv(16, dataset.num_classes) def forward(self, data): x, edge_index, batch = data.x, data.edge_index, data.batch x = self.conv1(x, edge_index) x = F.relu(x) #x = F.dropout(x, training=self.training) x = self.conv2(x, edge_index) x = F.relu(x) x = self.conv3(x, edge_index) x = torch_geometric.nn.global_max_pool(x, batch) return F.log_softmax(x, dim=1) loader = DataLoader(dataset, batch_size=32, shuffle=True) loader_test = DataLoader(dataset_test, batch_size=1, shuffle=False) for batch in loader: batch print(batch) print(batch.y[0]) print(np.array(batch.y)) break device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = Net().to(device) #data = dataset[0].to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=5e-4) model.train() n_train_correct = 0 n_train_count = 0 for epoch in range(200): for batch in loader: dat = batch.to(device) optimizer.zero_grad() out = model(dat) loss = F.nll_loss(out, dat.y) loss.backward() optimizer.step() n_train_correct += (torch.max(out, 1)[1].view( batch.y.size()) == batch.y).sum().item() n_train_count += 32 #batch.y.size() dev_acc = 100. * n_train_correct / n_train_count print('Train accuracy: ', dev_acc) # calculate accuracy on validation set n_dev_correct = 0 n_count = 0 with torch.no_grad(): for batch in loader_test: data = batch.to(device) answer = model(data) n_dev_correct += (torch.max(answer, 1)[1].view( batch.y.size()) == batch.y).sum().item() n_count += 1 dev_acc = 100. * n_dev_correct / n_count print('Test accuracy: ', dev_acc)
import os.path as osp import torch import torch.nn.functional as F from torch_geometric.datasets import MNISTSuperpixels from torch_geometric.data import DataListLoader import torch_geometric.transforms as T from torch_geometric.nn import SplineConv, global_mean_pool, DataParallel dataset = MNISTSuperpixels("/git/data/GraphData/MNIST", transform=T.Cartesian()).shuffle() loader = DataListLoader(dataset, batch_size=1024, shuffle=True) class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = SplineConv(dataset.num_features, 32, dim=2, kernel_size=5) self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5) self.lin1 = torch.nn.Linear(64, 128) self.lin2 = torch.nn.Linear(128, dataset.num_classes) def forward(self, data): print('Inside Model: num graphs: {}, device: {}'.format( data.num_graphs, data.batch.device)) x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr for i in range(3): x = F.elu(self.conv1(x, edge_index, edge_attr)) x = F.elu(self.conv2(x, edge_index, edge_attr)) x = global_mean_pool(x, data.batch)
# get the current node index of G idx = list(G.nodes()) # set the node sizes using node features size = x[idx] * 500 + 200 # set the node colors using node features color = [] for i in idx: grey = x[i] if grey == 0: color.append('skyblue') else: color.append('red') nx.draw(G, with_labels=True, node_size=size, node_color=color, pos=pos) plt.title("MNIST Superpixel") if __name__ == '__main__': path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST') image_dataset = MNIST(root=path, download=True) graph_dataset = MNISTSuperpixels(root=path) example = 25 image, label = image_dataset[example] data = graph_dataset[example] print("Image Label:", label) visualize(image, data)
def TrainingNet(dataset, modelName, params, num_pre_epochs, num_epochs, NumCutoff, optimizerName, MonteSize, savepath): Batch_size = int(params[0]) root = '/git/data/GraphData/' + dataset TestAccs = [] for Monte_iter in range(MonteSize): # Data start_epoch = 0 # start from epoch 0 or last checkpoint epoch NewNetworkSizeAdjust = [] WeightsDynamicsEvolution = [] # model if dataset == 'Cora' or dataset == 'Citeseer' or dataset == 'Pubmed': datasetroot = Planetoid(root=root, name=dataset, transform=T.NormalizeFeatures()).shuffle() trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) [net, model_to_save] = ModelAndSave(dataset, modelName, datasetroot, params, num_epochs) criterion = nn.CrossEntropyLoss() elif dataset == "CoraFull": datasetroot = CoraFull(root=root, transform=T.NormalizeFeatures()).shuffle() trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) [net, model_to_save] = ModelAndSave(dataset, modelName, datasetroot, params, num_epochs) elif dataset == "Amazon": datasetroot = Amazon(root, "Photo", transform=None, pre_transform=None) trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) [net, model_to_save] = ModelAndSave(dataset, modelName, datasetroot, params, num_epochs) elif dataset == 'ENZYMES' or dataset == 'MUTAG': datasetroot = TUDataset(root, name=dataset, use_node_attr=True) Num = len(datasetroot) // 10 global train_dataset, test_dataset train_dataset = datasetroot[:Num] test_dataset = datasetroot[Num:] trainloader = DataLoader(train_dataset, batch_size=Batch_size) testloader = DataLoader(test_dataset, batch_size=60) [net, model_to_save] = ModelAndSave(dataset, modelName, datasetroot, params, num_epochs) elif dataset == "PPI": train_dataset = PPI(root, split='train') test_dataset = PPI(root, split='test') trainloader = DataLoader(train_dataset, batch_size=Batch_size, shuffle=True) testloader = DataLoader(test_dataset, batch_size=1, shuffle=False) [net, model_to_save] = ModelAndSave(dataset, modelName, train_dataset, params, num_epochs) criterion = torch.nn.BCEWithLogitsLoss() elif dataset == "Reddit": datasetroot = Reddit(root) trainloader = DataListLoader(datasetroot, batch_size=1, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=2, shuffle=False) [net, model_to_save] = ModelAndSave(dataset, modelName, datasetroot, params, num_epochs) criterion = torch.nn.BCEWithLogitsLoss() elif dataset == 'MNIST': datasetroot = MNISTSuperpixels(root=root, transform=T.Cartesian()) trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True) testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False) [net, model_to_save] = ModelAndSave(dataset, modelName, datasetroot, params, num_epochs) elif dataset == 'CIFAR10': pass else: raise Exception("Input wrong datatset!!") FileName = "{}-{}-param_{}_{}_{}_{}-monte_{}".format( dataset, modelName, params[0], params[1], params[2], params[3], Monte_iter) print('Let\'s use', torch.cuda.device_count(), 'GPUs!') global device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') optimizer = optim.Adam(net.parameters(), lr=params[3], betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False) criterion = nn.CrossEntropyLoss() net = net.to(device) #cudnn.benchmark = True logging( 'dataset:{}, Batch size: {}, Number of layers:{} ConCoeff: {}, LR:{}, MonteSize:{}' .format(dataset, params[0], params[1], params[2], params[3], Monte_iter)) mark = "{}/{}Convergence/DiagElement-{}".format( savepath, dataset, FileName) PreTrainConvergence, PreTestConvergence, PreTestAcc = TrainPart( modelName, datasetroot, start_epoch, num_pre_epochs, trainloader, testloader, net, optimizer, criterion, NumCutoff, mark, False, model_to_save) print( 'dataset: {}, model name: {}, Number epoches: {}, Pre-train error is: {}, Pre-test error is: {}, test acc is {}' .format(dataset, modelName, num_pre_epochs, PreTrainConvergence[-1], PreTestConvergence[-1], PreTestAcc[-1])) NewNetworksize, NewNetworkWeight = RetainNetworkSize(net, params[2])[0:2] NetworkInfo = [NewNetworksize[0:-1], NewNetworkWeight] OptimizedNet = ChooseModel(modelName, datasetroot, NetworkInfo) NewNetworksize.insert(0, datasetroot.num_features) NewNetworkSizeAdjust.append(NewNetworksize[0:-1]) print(NewNetworkSizeAdjust) #OptimizedNet.apply(init_weights) #OptimizedNet = DataParallel(OptimizedNet) OptimizedNet = OptimizedNet.to(device) cudnn.benchmark = True criterionNew = nn.CrossEntropyLoss() if optimizerName == "SGD": optimizerNew = getattr(optim, optimizerName)(OptimizedNet.parameters(), lr=params[3], momentum=0.9, weight_decay=5e-4) elif optimizerName == "Adam": optimizerNew = getattr(optim, optimizerName)(OptimizedNet.parameters(), lr=params[3], betas=(0.9, 0.999), eps=1e-08, weight_decay=5e-4, amsgrad=False) TrainConvergence, TestConvergence, TestAcc = TrainPart( modelName, datasetroot, start_epoch, num_epochs, trainloader, testloader, OptimizedNet, optimizerNew, criterionNew, NumCutoff, mark, True, model_to_save) np.save( "{}/{}Convergence/TrainConvergence-{}".format( savepath, dataset, FileName), TrainConvergence) np.save( "{}/{}Convergence/TestConvergence-{}".format( savepath, dataset, FileName), TestConvergence) np.save( "{}/{}Convergence/NewNetworkSizeAdjust-{}".format( savepath, dataset, FileName), NewNetworkSizeAdjust) #np.save(savepath+'TestConvergence-'+FileName,TestConvergence) #torch.cuda.empty_cache() print( 'dataset: {}, model name:{}, resized network size is {}, Number epoches:{}, Train error is: {}, Test error is: {}, test acc is {}\n' .format(dataset, modelName, NewNetworksize[0:-1], num_epochs, TrainConvergence[-1], TestConvergence[-1], TestAcc[-1])) TestAccs.append(TestAcc) np.save( "{}/{}Convergence/MeanTestAccs-{}".format(savepath, dataset, FileName), TestAccs) print("The change of test error is:{}".format(TestAccs)) print_nvidia_useage()
num = 3 model_path = dir_path + "/cmodels/12_global_edge_attr_test/C_300.pt" dataset_path = dir_path + '/dataset/cartesian/' dataset_path += str(num) + "s/" if num != -1 else "all_nums" def pf(data): return data.y == num pre_filter = pf if num != -1 else None train_dataset = MNISTSuperpixels(dataset_path, True, pre_transform=T.Cartesian(), pre_filter=pre_filter) train_loader = DataLoader(train_dataset, batch_size=128) print("loaded data") pretrained_model = torch.load(model_path, map_location=device) C = MoNet(25) C.load_state_dict(pretrained_model.state_dict()) torch.save(pretrained_model.state_dict(), "../mnist_superpixels/evaluation/C_state_dict.pt") print("loaded model)") # Testing Classification - remember to comment the first return in MoNet.forward()
def main(args): args = init_dirs(args) pt = T.Cartesian() if args.cartesian else T.Polar() if args.dataset == 'sp': train_dataset = MNISTSuperpixels(args.dataset_path, True, pre_transform=pt) test_dataset = MNISTSuperpixels(args.dataset_path, False, pre_transform=pt) train_loader = tgDataLoader(train_dataset, batch_size=args.batch_size, shuffle=True) test_loader = tgDataLoader(test_dataset, batch_size=args.batch_size) elif args.dataset == 'sm': train_dataset = MNISTGraphDataset(args.dataset_path, args.num_hits, train=True) train_loader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch_size, pin_memory=True) test_dataset = MNISTGraphDataset(args.dataset_path, args.num_hits, train=False) test_loader = DataLoader(test_dataset, shuffle=True, batch_size=args.batch_size, pin_memory=True) if(args.load_model): C = torch.load(args.model_path + args.name + "/C_" + str(args.start_epoch) + ".pt").to(device) else: C = MoNet(args.kernel_size).to(device) C_optimizer = torch.optim.Adam(C.parameters(), lr=args.lr, weight_decay=args.weight_decay) if(args.scheduler): C_scheduler = torch.optim.lr_scheduler.StepLR(C_optimizer, args.decay_step, gamma=args.lr_decay) train_losses = [] test_losses = [] def plot_losses(epoch, train_losses, test_losses): fig = plt.figure() ax1 = fig.add_subplot(1, 2, 1) ax1.plot(train_losses) ax1.set_title('training') ax2 = fig.add_subplot(1, 2, 2) ax2.plot(test_losses) ax2.set_title('testing') plt.savefig(args.losses_path + args.name + "/" + str(epoch) + ".png") plt.close() def save_model(epoch): torch.save(C, args.model_path + args.name + "/C_" + str(epoch) + ".pt") def train_C(data, y): C.train() C_optimizer.zero_grad() output = C(data) # nll_loss takes class labels as target, so one-hot encoding is not needed C_loss = F.nll_loss(output, y) C_loss.backward() C_optimizer.step() return C_loss.item() def test(epoch): C.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data in test_loader: if args.dataset == 'sp': output = C(data.to(device)) y = data.y.to(device) elif args.dataset == 'sm': output = C(tg_transform(args, data[0].to(device))) y = data[1].to(device) test_loss += F.nll_loss(output, y, size_average=False).item() pred = output.data.max(1, keepdim=True)[1] correct += pred.eq(y.data.view_as(pred)).sum() test_loss /= len(test_loader.dataset) test_losses.append(test_loss) print('test') f = open(args.out_path + args.name + '.txt', 'a') print(args.out_path + args.name + '.txt') s = "After {} epochs, on test set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format(epoch, test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset)) print(s) f.write(s) f.close() for i in range(args.start_epoch, args.num_epochs): print("Epoch %d %s" % ((i + 1), args.name)) C_loss = 0 test(i) for batch_ndx, data in tqdm(enumerate(train_loader), total=len(train_loader)): if args.dataset == 'sp': C_loss += train_C(data.to(device), data.y.to(device)) elif args.dataset == 'sm': C_loss += train_C(tg_transform(args, data[0].to(device)), data[1].to(device)) train_losses.append(C_loss / len(train_loader)) if(args.scheduler): C_scheduler.step() if((i + 1) % 10 == 0): save_model(i + 1) plot_losses(i + 1, train_losses, test_losses) test(args.num_epochs)
import os.path as osp import torch import torch.nn as nn import torch.nn.functional as F from torch_geometric.datasets import MNISTSuperpixels import torch_geometric.transforms as T from torch_geometric.data import DataLoader from torch_geometric.utils import normalized_cut from torch_geometric.nn import (NNConv, graclus, max_pool, max_pool_x, global_mean_pool) path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST') train_dataset = MNISTSuperpixels(path, True, transform=T.Cartesian()) test_dataset = MNISTSuperpixels(path, False, transform=T.Cartesian()) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=64) d = train_dataset def normalized_cut_2d(edge_index, pos): row, col = edge_index edge_attr = torch.norm(pos[row] - pos[col], p=2, dim=1) return normalized_cut(edge_index, edge_attr, num_nodes=pos.size(0)) class Net(nn.Module): def __init__(self): super(Net, self).__init__() n1 = nn.Sequential(nn.Linear(2, 25), nn.ReLU(), nn.Linear(25, 32)) self.conv1 = NNConv(d.num_features, 32, n1)
def get_dataset(name, sparse=True, feat_str="deg+ak3+reall", root=None, aug=None, aug_ratio=None): if root is None or root == '': path = osp.join(osp.expanduser('~'), 'pyG_data', name) else: path = osp.join(root, name) degree = feat_str.find("deg") >= 0 onehot_maxdeg = re.findall("odeg(\d+)", feat_str) onehot_maxdeg = int(onehot_maxdeg[0]) if onehot_maxdeg else None k = re.findall("an{0,1}k(\d+)", feat_str) k = int(k[0]) if k else 0 groupd = re.findall("groupd(\d+)", feat_str) groupd = int(groupd[0]) if groupd else 0 remove_edges = re.findall("re(\w+)", feat_str) remove_edges = remove_edges[0] if remove_edges else 'none' edge_noises_add = re.findall("randa([\d\.]+)", feat_str) edge_noises_add = float(edge_noises_add[0]) if edge_noises_add else 0 edge_noises_delete = re.findall("randd([\d\.]+)", feat_str) edge_noises_delete = float( edge_noises_delete[0]) if edge_noises_delete else 0 centrality = feat_str.find("cent") >= 0 coord = feat_str.find("coord") >= 0 pre_transform = FeatureExpander(degree=degree, onehot_maxdeg=onehot_maxdeg, AK=k, centrality=centrality, remove_edges=remove_edges, edge_noises_add=edge_noises_add, edge_noises_delete=edge_noises_delete, group_degree=groupd).transform print(aug, aug_ratio) if 'MNIST' in name or 'CIFAR' in name: if name == 'MNIST_SUPERPIXEL': train_dataset = MNISTSuperpixels(path, True, pre_transform=pre_transform, transform=T.Cartesian()) test_dataset = MNISTSuperpixels(path, False, pre_transform=pre_transform, transform=T.Cartesian()) else: train_dataset = ImageDataset(path, name, True, pre_transform=pre_transform, coord=coord, processed_file_prefix="data_%s" % feat_str) test_dataset = ImageDataset(path, name, False, pre_transform=pre_transform, coord=coord, processed_file_prefix="data_%s" % feat_str) dataset = (train_dataset, test_dataset) else: dataset = TUDatasetExt(path, name, pre_transform=pre_transform, use_node_attr=True, processed_filename="data_%s.pt" % feat_str, aug=aug, aug_ratio=aug_ratio) dataset.data.edge_attr = None return dataset
import os.path as osp import torch import torch.nn.functional as F from torch_geometric.datasets import MNISTSuperpixels from torch_geometric.data import DataListLoader import torch_geometric.transforms as T from torch_geometric.nn import SplineConv, global_mean_pool, DataParallel path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data', 'MNIST') dataset = MNISTSuperpixels(path, transform=T.Cartesian()).shuffle() loader = DataListLoader(dataset, batch_size=1024, shuffle=True) class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = SplineConv(dataset.num_features, 32, dim=2, kernel_size=5) self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5) self.lin1 = torch.nn.Linear(64, 128) self.lin2 = torch.nn.Linear(128, dataset.num_classes) def forward(self, data): print('Inside Model: num graphs: {}, device: {}'.format( data.num_graphs, data.batch.device)) x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr x = F.elu(self.conv1(x, edge_index, edge_attr)) x = F.elu(self.conv2(x, edge_index, edge_attr)) x = global_mean_pool(x, data.batch) x = F.elu(self.lin1(x))
start_time = time.time() test_acc = validate(model, test_loader, device) print('Test, Time: {:.4f}, Accuracy: {:.4f}'\ .format(time.time() - start_time, test_acc), file=file) return model if __name__ == '__main__': # SuperPixels dataset sp_path = os.path.join(os.path.dirname(os.path.realpath("/")), "MNISTSuperpixel") sp_test_dataset = MNISTSuperpixels(sp_path, train=False, transform=T.Cartesian()) sp_train_loader_25, sp_val_loader_25 = get_train_val_loader( sp_path, train_batch_size=25, val_batch_size=25) sp_train_loader_64, sp_val_loader_64 = get_train_val_loader( sp_path, train_batch_size=64, val_batch_size=64) sp_test_loader = DataLoader(sp_test_dataset, batch_size=1, shuffle=False) # Skeletons dataset sk_path = "dataset" sk_train_dataset = MNISTSkeleton(sk_path, "train", transform=T.Polar()) sk_test_dataset = MNISTSkeleton(sk_path, "test", transform=T.Polar()) sk_val_dataset = MNISTSkeleton(sk_path, "val", transform=T.Polar()) sk_train_loader = DataLoader(sk_train_dataset, batch_size=64, shuffle=True) sk_val_loader = DataLoader(sk_val_dataset, batch_size=64, shuffle=False)
parser.add_argument('--weight_decay', type=float, default=5e-4) parser.add_argument('--batch_size', type=int, default=64) parser.add_argument('--epochs', type=int, default=300) parser.add_argument('--seed', type=int, default=1) args = parser.parse_args() args.data_fp = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', args.dataset) device = torch.device('cuda', args.device_idx) # deterministic torch.manual_seed(args.seed) cudnn.benchmark = False cudnn.deterministic = True train_dataset = MNISTSuperpixels(args.data_fp, True, pre_transform=T.Polar()) test_dataset = MNISTSuperpixels(args.data_fp, False, pre_transform=T.Polar()) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=64) def normalized_cut_2d(edge_index, pos): row, col = edge_index edge_attr = torch.norm(pos[row] - pos[col], p=2, dim=1) return normalized_cut(edge_index, edge_attr, num_nodes=pos.size(0)) class MoNet(torch.nn.Module): def __init__(self, kernel_size): super(MoNet, self).__init__() self.conv1 = GMMConv(1, 32, dim=2, kernel_size=kernel_size)
def main(args): torch.manual_seed(4) torch.autograd.set_detect_anomaly(True) name = [args.name] if WGAN: name.append('wgan') if GRU: name.append('gru') if GCNN: name.append('gcnn') # name.append('num_iters_{}'.format(args.num_iters)) # name.append('num_critic_{}'.format(args.num_critic)) args.name = '_'.join(name) args.model_path = args.dir_path + '/models/' args.losses_path = args.dir_path + '/losses/' args.args_path = args.dir_path + '/args/' args.figs_path = args.dir_path + '/figs/' args.dataset_path = args.dir_path + '/dataset/' args.err_path = args.dir_path + '/err/' if(not exists(args.model_path)): mkdir(args.model_path) if(not exists(args.losses_path)): mkdir(args.losses_path) if(not exists(args.args_path)): mkdir(args.args_path) if(not exists(args.figs_path)): mkdir(args.figs_path) if(not exists(args.err_path)): mkdir(args.err_path) if(not exists(args.dataset_path)): mkdir(args.dataset_path) try: # python2 file_tmp = urllib.urlretrieve(url, filename=None)[0] except: # python3 file_tmp = urllib.request.urlretrieve(url, filename=args.dataset)[0] prev_models = [f[:-4] for f in listdir(args.args_path)] # removing .txt if (args.name in prev_models): print("name already used") # if(not args.load_model): # sys.exit() else: mkdir(args.losses_path + args.name) mkdir(args.model_path + args.name) mkdir(args.figs_path + args.name) if(not args.load_model): f = open(args.args_path + args.name + ".txt", "w+") f.write(str(vars(args))) f.close() else: f = open(args.args_path + args.name + ".txt", "r") temp = args.start_epoch args = eval(f.read()) f.close() args.load_model = True args.start_epoch = temp # return args2 def pf(data): return data.y == args.num pre_filter = pf if args.num != -1 else None print("loading") # Change to True !! X = SuperpixelsDataset(args.dataset_path, args.num_hits, train=TRAIN, num=NUM, device=device) tgX = MNISTSuperpixels(args.dir_path, train=TRAIN, pre_transform=T.Cartesian(), pre_filter=pre_filter) X_loaded = DataLoader(X, shuffle=True, batch_size=args.batch_size, pin_memory=True) tgX_loaded = tgDataLoader(tgX, shuffle=True, batch_size=args.batch_size) print("loaded") if(args.load_model): G = torch.load(args.model_path + args.name + "/G_" + str(args.start_epoch) + ".pt") D = torch.load(args.model_path + args.name + "/D_" + str(args.start_epoch) + ".pt") else: G = Graph_Generator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.gru_hidden_size, args.gru_num_layers, args.num_iters, args.num_hits, args.dropout, args.leaky_relu_alpha, hidden_node_size=args.hidden_node_size, int_diffs=INT_DIFFS, gru=GRU, device=device).to(device) if(GCNN): D = MoNet(kernel_size=args.kernel_size, dropout=args.dropout, device=device).to(device) # D = Gaussian_Discriminator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.gru_hidden_size, args.gru_num_layers, args.num_iters, args.num_hits, args.dropout, args.leaky_relu_alpha, kernel_size=args.kernel_size, hidden_node_size=args.hidden_node_size, int_diffs=INT_DIFFS, gru=GRU).to(device) else: D = Graph_Discriminator(args.node_feat_size, args.fe_hidden_size, args.fe_out_size, args.gru_hidden_size, args.gru_num_layers, args.num_iters, args.num_hits, args.dropout, args.leaky_relu_alpha, hidden_node_size=args.hidden_node_size, int_diffs=INT_DIFFS, gru=GRU, device=device).to(device) print("Models loaded") if(WGAN): G_optimizer = optim.RMSprop(G.parameters(), lr=args.lr_gen) D_optimizer = optim.RMSprop(D.parameters(), lr=args.lr_disc) else: G_optimizer = optim.Adam(G.parameters(), lr=args.lr_gen, weight_decay=5e-4) D_optimizer = optim.Adam(D.parameters(), lr=args.lr_disc, weight_decay=5e-4) print("optimizers loaded") normal_dist = Normal(0, 0.2) def wasserstein_loss(y_out, y_true): return -torch.mean(y_out * y_true) if(WGAN): criterion = wasserstein_loss else: if(LSGAN): criterion = torch.nn.MSELoss() else: criterion = torch.nn.BCELoss() # print(criterion(torch.tensor([1.0]),torch.tensor([-1.0]))) def gen(num_samples, noise=0): if(noise == 0): noise = normal_dist.sample((num_samples, args.num_hits, args.hidden_node_size)).to(device) return G(noise) # transform my format to torch_geometric's def tg_transform(X): batch_size = X.size(0) cutoff = 0.32178 # found empirically to match closest to Superpixels pos = X[:, :, :2] x1 = pos.repeat(1, 1, 75).reshape(batch_size, 75*75, 2) x2 = pos.repeat(1, 75, 1) diff_norms = torch.norm(x2 - x1 + 1e-12, dim=2) diff = x2-x1 diff = diff[diff_norms < cutoff] norms = diff_norms.reshape(batch_size, 75, 75) neighborhood = torch.nonzero(norms < cutoff, as_tuple=False) edge_attr = diff[neighborhood[:, 1] != neighborhood[:, 2]] neighborhood = neighborhood[neighborhood[:, 1] != neighborhood[:, 2]] # remove self-loops unique, counts = torch.unique(neighborhood[:, 0], return_counts=True) edge_slices = torch.cat((torch.tensor([0]).to(device), counts.cumsum(0))) edge_index = neighborhood[:, 1:].transpose(0, 1) # normalizing edge attributes edge_attr_list = list() for i in range(batch_size): start_index = edge_slices[i] end_index = edge_slices[i+1] temp = diff[start_index:end_index] max = torch.max(temp) temp = temp/(2*max + 1e-12) + 0.5 edge_attr_list.append(temp) edge_attr = torch.cat(edge_attr_list) x = X[:, :, 2].reshape(batch_size*75, 1)+0.5 pos = 27*pos.reshape(batch_size*75, 2)+13.5 zeros = torch.zeros(batch_size*75, dtype=int).to(device) zeros[torch.arange(batch_size)*75] = 1 batch = torch.cumsum(zeros, 0)-1 return Batch(batch=batch, x=x, edge_index=edge_index.contiguous(), edge_attr=edge_attr, y=None, pos=pos) def draw_graph(graph, node_r, im_px): imd = im_px + node_r img = np.zeros((imd, imd), dtype=np.float) circles = [] for node in graph: circles.append((draw.circle_perimeter(int(node[1]), int(node[0]), node_r), draw.disk((int(node[1]), int(node[0])), node_r), node[2])) for circle in circles: img[circle[1]] = circle[2] return img def save_sample_outputs(name, epoch, dlosses, glosses): print("drawing figs") fig = plt.figure(figsize=(10, 10)) num_ims = 100 node_r = 30 im_px = 1000 gen_out = gen(args.batch_size).cpu().detach().numpy() for i in range(int(num_ims/args.batch_size)): gen_out = np.concatenate((gen_out, gen(args.batch_size).cpu().detach().numpy()), 0) gen_out = gen_out[:num_ims] gen_out[gen_out > 0.47] = 0.47 gen_out[gen_out < -0.5] = -0.5 gen_out = gen_out*[im_px, im_px, 1] + [(im_px+node_r)/2, (im_px+node_r)/2, 0.55] for i in range(1, num_ims+1): fig.add_subplot(10, 10, i) im_disp = draw_graph(gen_out[i-1], node_r, im_px) plt.imshow(im_disp, cmap=cm.gray_r, interpolation='nearest') plt.axis('off') plt.savefig(args.figs_path + args.name + "/" + str(epoch) + ".png") plt.close() plt.figure() plt.plot(dlosses, label='Discriminitive loss') plt.plot(glosses, label='Generative loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.legend() plt.savefig(args.losses_path + args.name + "/" + str(epoch) + ".png") plt.close() print("saved figs") def save_models(name, epoch): torch.save(G, args.model_path + args.name + "/G_" + str(epoch) + ".pt") torch.save(D, args.model_path + args.name + "/D_" + str(epoch) + ".pt") # from https://github.com/EmilienDupont/wgan-gp def gradient_penalty(real_data, generated_data): batch_size = real_data.size()[0] # Calculate interpolation alpha = torch.rand(batch_size, 1, 1) alpha = alpha.expand_as(real_data).to(device) interpolated = alpha * real_data.data + (1 - alpha) * generated_data.data interpolated = Variable(interpolated, requires_grad=True).to(device) del alpha torch.cuda.empty_cache() # Calculate probability of interpolated examples prob_interpolated = D(interpolated) # Calculate gradients of probabilities with respect to examples gradients = torch_grad(outputs=prob_interpolated, inputs=interpolated, grad_outputs=torch.ones(prob_interpolated.size()).to(device), create_graph=True, retain_graph=True, allow_unused=True)[0].to(device) gradients = gradients.contiguous() # Gradients have shape (batch_size, num_channels, img_width, img_height), # so flatten to easily take norm per example in batch gradients = gradients.view(batch_size, -1) # Derivatives of the gradient close to 0 can cause problems because of # the square root, so manually calculate norm and add epsilon gradients_norm = torch.sqrt(torch.sum(gradients ** 2, dim=1) + 1e-12) # Return gradient penalty return args.gp_weight * ((gradients_norm - 1) ** 2).mean() def train_D(data): D.train() D_optimizer.zero_grad() run_batch_size = data.shape[0] if not GCNN else data.y.shape[0] if(not WGAN): Y_real = torch.ones(run_batch_size, 1).to(device) Y_fake = torch.zeros(run_batch_size, 1).to(device) try: D_real_output = D(data) gen_ims = gen(run_batch_size) tg_gen_ims = tg_transform(gen_ims) use_gen_ims = tg_gen_ims if GCNN else gen_ims D_fake_output = D(use_gen_ims) if(WGAN): D_loss = D_fake_output.mean() - D_real_output.mean() + gradient_penalty(data, use_gen_ims) else: D_real_loss = criterion(D_real_output, Y_real) D_fake_loss = criterion(D_fake_output, Y_fake) D_loss = D_real_loss + D_fake_loss D_loss.backward() D_optimizer.step() except: print("Generated Images") print(gen_ims) print("Transformed Images") print(tg_gen_ims) print("Discriminator Output") print(D_fake_output) torch.save(gen_ims, args.err_path + args.name + "_gen_ims.pt") torch.save(tg_gen_ims.x, args.err_path + args.name + "_x.pt") torch.save(tg_gen_ims.pos, args.err_path + args.name + "_pos.pt") torch.save(tg_gen_ims.edge_index, args.err_path + args.name + "_edge_index.pt") return return D_loss.item() def train_G(): G.train() G_optimizer.zero_grad() if(not WGAN): Y_real = torch.ones(args.batch_size, 1).to(device) gen_ims = gen(args.batch_size) if(GCNN): gen_ims = tg_transform(gen_ims) D_fake_output = D(gen_ims) if(WGAN): G_loss = -D_fake_output.mean() else: G_loss = criterion(D_fake_output, Y_real) G_loss.backward() G_optimizer.step() return G_loss.item() D_losses = [] G_losses = [] # save_models(name, 0) # save_sample_outputs(args.name, 0, D_losses, G_losses) # @profile def train(): for i in range(args.start_epoch, args.num_epochs): print("Epoch %d %s" % ((i+1), args.name)) D_loss = 0 G_loss = 0 loader = tgX_loaded if GCNN else X_loaded for batch_ndx, data in tqdm(enumerate(loader), total=len(loader)): if(batch_ndx > 0 and batch_ndx % (args.num_critic+1) == 0): G_loss += train_G() else: D_loss += train_D(data.to(device)) if GCNN else train_D(data[0].to(device)) D_losses.append(D_loss/len(X_loaded)/2) G_losses.append(G_loss/len(X_loaded)) if((i+1) % 5 == 0): save_sample_outputs(args.name, i+1, D_losses, G_losses) if((i+1) % 5 == 0): save_models(args.name, i+1) train()