def main(dataset, times): adj, features, labels, idx_train, idx_val, idx_test = load_data(dataset) features = features.to(device) adj = adj.to(device) labels = labels.to(device) idx_train = idx_train.to(device) idx_val = idx_val.to(device) idx_test = idx_test.to(device) nclass = labels.max().item() + 1 acc_lst = list() for seed in random.sample(range(0, 100000), times): np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed(seed) # Model and optimizer # weight_decay 权值衰减,防止过拟合,调节模型复杂度对损失函数的影响 model = GCN(nfeat=features.shape[1], nhid=args.hidden, nclass=nclass, dropout=args.dropout) optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) model.to(device) # Train model t_total = time.time() for epoch in range(args.epochs): train(epoch, model, optimizer, adj, features, labels, idx_train, idx_val) print(f"Total time elapsed: {time.time() - t_total:.4f}s") # Testing acc_lst.append(test(model, adj, features, labels, idx_test)) print(acc_lst) print(np.mean(acc_lst))
def main(): config = get_args() dataset = Data(config) num_features = dataset.features.shape[1] num_classes = dataset.labels.max().item() + 1 model = GCN(config=config, num_features=num_features, num_classes=num_classes) solver = Solver(config, model, dataset) if torch.cuda.is_available(): model = model.to('cuda') criterion, best_model = solver.train() solver.test(criterion, best_model)
def main(args): start_time_str = time.strftime("_%m_%d_%H_%M_%S", time.localtime()) log_path = os.path.join(args.log_dir, args.model + start_time_str) if not os.path.exists(log_path): os.mkdir(log_path) logging.basicConfig(filename=os.path.join(log_path, 'log_file'), filemode='w', format='| %(asctime)s |\n%(message)s', datefmt='%b %d %H:%M:%S', level=logging.INFO) logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) logging.info(args) # load data if args.model in ['EGNN', 'ECConv', 'GTEA-ST']: data = Dataset(data_dir=args.data_dir, batch_size=args.batch_size, use_static=True) else: data = Dataset(data_dir=args.data_dir, batch_size=args.batch_size) g = data.g # features = torch.FloatTensor(data.features) # labels = torch.LongTensor(data.labels) train_loader = data.train_loader val_loader = data.val_loader test_loader = data.test_loader num_nodes = data.num_nodes node_in_dim = args.node_in_dim num_edges = data.num_edges edge_in_dim = data.edge_in_dim edge_timestep_len = data.edge_timestep_len num_train_samples = data.num_train_samples num_val_samples = data.num_val_samples num_test_samples = data.num_test_samples logging.info("""----Data statistics------' #Nodes %d #Edges %d #Node_feat %d #Edge_feat %d #Edge_timestep %d #Train samples %d #Val samples %d #Test samples %d""" % (num_nodes, num_edges, node_in_dim, edge_in_dim, edge_timestep_len, num_train_samples, num_val_samples, num_test_samples)) device = torch.device("cuda:"+str(args.gpu) if torch.cuda.is_available() and args.gpu >=0 else "cpu") infer_device = device if args.infer_gpu else torch.device('cpu') # g = g.to(device) # create model if args.model == 'GCN': model = GCN(num_nodes=num_nodes, in_feats=node_in_dim, n_hidden=args.node_hidden_dim, n_layers=args.num_layers, activation=F.relu, dropout=args.dropout) elif args.model == 'GraphSAGE': model = GraphSAGE(num_nodes=num_nodes, in_feats=node_in_dim, n_hidden=args.node_hidden_dim, n_layers=args.num_layers, activation=F.relu, dropout=args.dropout) elif args.model == 'GAT': model = GAT(num_nodes=num_nodes, in_dim=node_in_dim, hidden_dim=args.node_hidden_dim, num_layers=args.num_layers, num_heads=args.num_heads) elif args.model == 'ECConv': model = ECConv(num_nodes=num_nodes, node_in_dim=node_in_dim, edge_in_dim=edge_in_dim, hidden_dim=args.node_hidden_dim, num_layers=args.num_layers, drop_prob=args.dropout, device=device) elif args.model == 'EGNN': model = EGNN(num_nodes=num_nodes, node_in_dim=node_in_dim, edge_in_dim=edge_in_dim, hidden_dim=args.node_hidden_dim, num_layers=args.num_layers, drop_prob=args.dropout, device=device) elif args.model == 'GTEA-ST': model = GTEAST(num_nodes=num_nodes, node_in_dim=node_in_dim, edge_in_dim=edge_in_dim, node_hidden_dim=args.node_hidden_dim, num_layers=args.num_layers, drop_prob=args.dropout, device=device) elif args.model == 'TGAT': model = TGAT(num_nodes=num_nodes, node_in_dim=node_in_dim, node_hidden_dim=args.node_hidden_dim, edge_in_dim=edge_in_dim-1, time_hidden_dim=args.time_hidden_dim, num_class=0, num_layers=args.num_layers, num_heads=args.num_heads, device=device, drop_prob=args.dropout) elif args.model == 'GTEA-LSTM': model = GTEALSTM(num_nodes=num_nodes, node_in_dim=node_in_dim, node_hidden_dim=args.node_hidden_dim, edge_in_dim=edge_in_dim, num_class=0, num_layers=args.num_layers, num_time_layers=args.num_lstm_layers, bidirectional=args.bidirectional, device=device, drop_prob=args.dropout) elif args.model == 'GTEA-LSTM+T2V': model = GTEALSTMT2V(num_nodes=num_nodes, node_in_dim=node_in_dim, node_hidden_dim=args.node_hidden_dim, edge_in_dim=edge_in_dim-1, time_hidden_dim=args.time_hidden_dim, num_class=0, num_layers=args.num_layers, num_time_layers=args.num_lstm_layers, bidirectional=args.bidirectional, device=device, drop_prob=args.dropout) elif args.model == 'GTEA-Trans': model = GTEATrans(num_nodes=num_nodes, node_in_dim=node_in_dim, node_hidden_dim=args.node_hidden_dim, edge_in_dim=edge_in_dim, num_class=0, num_layers=args.num_layers, num_heads=args.num_heads, num_time_layers=args.num_lstm_layers, device=device, drop_prob=args.dropout) elif args.model == 'GTEA-Trans+T2V': model = GTEATransT2V(num_nodes=num_nodes, node_in_dim=node_in_dim, node_hidden_dim=args.node_hidden_dim, edge_in_dim=edge_in_dim-1, time_hidden_dim=args.time_hidden_dim, num_class=0, num_layers=args.num_layers, num_heads=args.num_heads, num_time_layers=args.num_lstm_layers, device=device, drop_prob=args.dropout) else: logging.info('Model {} not found.'.format(args.model)) exit(0) # send model to device model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) checkpoint_path = os.path.join(log_path, str(args.model) + '_checkpoint.pt') trainer = Trainer(g=g, model=model, optimizer=optimizer, epochs=args.epochs, train_loader=train_loader, val_loader=val_loader, test_loader=test_loader, patience=args.patience, batch_size=args.batch_size, num_neighbors=args.num_neighbors, num_layers=args.num_layers, num_workers=args.num_workers, device=device, infer_device=infer_device, log_path=log_path, checkpoint_path=checkpoint_path) logging.info('Start training') best_val_result, test_result = trainer.train() # recording the result line = [start_time_str[1:]] + [args.model] + ['K=' + str(args.use_K)] + \ [str(x) for x in best_val_result] + [str(x) for x in test_result] + [str(args)] line = ','.join(line) + '\n' with open(os.path.join(args.log_dir, str(args.model) + '_result.csv'), 'a') as f: f.write(line)
def main(): # Make dir temp = "./tmp" os.makedirs(temp, exist_ok=True) # Load data start = time.time() (train_adj, full_adj, train_feats, test_feats, y_train, y_val, y_test, train_mask, val_mask, test_mask, _, val_nodes, test_nodes, num_data, visible_data) = utils.load_data(args.dataset) print('Loaded data in {:.2f} seconds!'.format(time.time() - start)) start = time.time() # Prepare Train Data if args.batch_size > 1: start = time.time() _, parts = utils.partition_graph(train_adj, visible_data, args.num_clusters_train) print('Partition graph in {:.2f} seconds!'.format(time.time() - start)) parts = [np.array(pt) for pt in parts] else: start = time.time() (parts, features_batches, support_batches, y_train_batches, train_mask_batches) = utils.preprocess( train_adj, train_feats, y_train, train_mask, visible_data, args.num_clusters_train, diag_lambda=args.diag_lambda) print('Partition graph in {:.2f} seconds!'.format(time.time() - start)) # Prepare valid Data start = time.time() (_, val_features_batches, val_support_batches, y_val_batches, val_mask_batches) = utils.preprocess( full_adj, test_feats, y_val, val_mask, np.arange(num_data), args.num_clusters_val, diag_lambda=args.diag_lambda) print('Partition graph in {:.2f} seconds!'.format(time.time() - start)) # Prepare Test Data start = time.time() (_, test_features_batches, test_support_batches, y_test_batches, test_mask_batches) = utils.preprocess( full_adj, test_feats, y_test, test_mask, np.arange(num_data), args.num_clusters_test, diag_lambda=args.diag_lambda) print('Partition graph in {:.2f} seconds!'.format(time.time() - start)) idx_parts = list(range(len(parts))) # model model = GCN( fan_in=_in, fan_out=_out, layers=args.layers, dropout=args.dropout, normalize=True, bias=False, precalc=True).float() model.to(torch.device('cuda')) print(model) # Loss Function if args.multilabel: criterion = torch.nn.BCEWithLogitsLoss() else: criterion = torch.nn.CrossEntropyLoss() # Optimization Algorithm optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # Learning Rate Schedule # scheduler = torch.optim.lr_scheduler.OneCycleLR( # optimizer, max_lr=args.lr, steps_per_epoch=int(args.num_clusters_train/args.batch_size), epochs=args.epochs+1, # anneal_strategy='linear') pbar = tqdm.tqdm(total=args.epochs, dynamic_ncols=True) for epoch in range(args.epochs + 1): # Train np.random.shuffle(idx_parts) start = time.time() avg_loss = 0 total_correct = 0 n_nodes = 0 if args.batch_size > 1: (features_batches, support_batches, y_train_batches, train_mask_batches) = utils.preprocess_multicluster( train_adj, parts, train_feats, y_train, train_mask, args.num_clusters_train, args.batch_size, args.diag_lambda) for pid in range(len(features_batches)): # Use preprocessed batch data features_b = features_batches[pid] support_b = support_batches[pid] y_train_b = y_train_batches[pid] train_mask_b = train_mask_batches[pid] loss, pred, labels = train( model.train(), criterion, optimizer, features_b, support_b, y_train_b, train_mask_b, torch.device('cuda')) avg_loss += loss n_nodes += pred.squeeze().numel() total_correct += torch.eq(pred.squeeze(), labels.squeeze()).sum().item() else: np.random.shuffle(idx_parts) for pid in idx_parts: # use preprocessed batch data features_b = features_batches[pid] support_b = support_batches[pid] y_train_b = y_train_batches[pid] train_mask_b = train_mask_batches[pid] loss, pred, labels = train( model.train(), criterion, optimizer, features_b, support_b, y_train_b, train_mask_b, torch.device('cuda')) avg_loss = loss.item() n_nodes += pred.squeeze().numel() total_correct += torch.eq(pred.squeeze(), labels.squeeze()).sum().item() train_acc = total_correct / n_nodes # Write Train stats to tensorboard writer.add_scalar('time/train', time.time() - start, epoch) writer.add_scalar('loss/train', avg_loss/len(features_batches), epoch) writer.add_scalar('acc/train', train_acc, epoch) # Validation cost, acc, micro, macro = evaluate( model.eval(), criterion, val_features_batches, val_support_batches, y_val_batches, val_mask_batches, val_nodes, torch.device("cuda")) # Write Valid stats to tensorboard writer.add_scalar('acc/valid', acc, epoch) writer.add_scalar('mi_F1/valid', micro, epoch) writer.add_scalar('ma_F1/valid', macro, epoch) writer.add_scalar('loss/valid', cost, epoch) pbar.set_postfix({"t": avg_loss/len(features_batches),"t_acc": train_acc, "v": cost, "v_acc": acc}) pbar.update() pbar.close() # Test if args.test == 1: # Test on cpu cost, acc, micro, macro = test( model.eval(), criterion, test_features_batches, test_support_batches, y_test_batches, test_mask_batches, torch.device("cpu")) writer.add_scalar('acc/test', acc, epoch) writer.add_scalar('mi_F1/test', micro, epoch) writer.add_scalar('ma_F1/test', macro, epoch) writer.add_scalar('loss/test', cost, epoch) print('test: acc: {:.4f}'.format(acc)) print('test: mi_f1: {:.4f}, ma_f1: {:.4f}'.format(micro, macro))
def cross_validation(): ##=======================Load Data================================================ # Load data population = 'PTE' epidata = np.load(population + '_graphs_gcn.npz') adj_epi = torch.from_numpy(calc_DAD(epidata)).float().to( device) # n_subjects*16 *16 features_epi = torch.from_numpy(epidata['features']).float().to( device) # n_subjectsx16x171 # n_subjects = features_epi.shape[0] # num_train = int(n_subjects * args.rate) # train_adj_epi = adj_epi[:num_train, :, :] # train_features_epi = features_epi[:num_train, :, :] # test_adj_epi = adj_epi[num_train:, :, :] # test_features_epi = features_epi[num_train:, :, :] population = 'NONPTE' nonepidata = np.load(population + '_graphs_gcn.npz') adj_non = torch.from_numpy(calc_DAD(nonepidata)).float().to(device) features_non = torch.from_numpy(nonepidata['features']).float().to( device) #subjects x 16 x 171 # print("DAD shape:") # print(adj_non.shape, adj_epi.shape) ## for now we are using the same number of epi , non epi training samples. # n_subjects_non = features_non.shape[0] # num_train_non = int(n_subjects_non * args.rate) # train_adj_non = adj_non[:num_train_non, :, :] # train_features_non = features_non[:num_train_non, :, :] # test_adj_non = adj_non[num_train_non:, :, :] # test_features_non = features_non[num_train_non:, :, :] features = torch.cat([features_epi, features_non]) adj = torch.cat([adj_epi, adj_non]) labels = torch.from_numpy( np.hstack((np.ones(adj_epi.shape[0]), np.zeros(adj_non.shape[0])))).long().to(device) print(features.shape, adj.shape, labels.shape) iterations = args.iters acc_iter = [] auc_iter = [] for i in range(iterations): kfold = StratifiedKFold(n_splits=36, shuffle=True) # the folds are made by preserving the percentage of samples for each class. acc = [] max_epochs = [] test_true = [] probs_fold = [] # epochs_choices = [] features_numpy = features.cpu().numpy() labels_numpy = labels.cpu().numpy() adj_numpy = adj.cpu().numpy() for train_ind, test_ind in kfold.split(features_numpy, labels_numpy): # Model and optimizer model = GCN( nfeat=features_epi.shape[2] // args.ngroup, nhid=[200, 200, 50], # nhid=[200, 200, 100, 50], nclass=2, #labels.max().item() + 1, dropout=args.dropout) optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # optimizer = optim.RMSprop(model.parameters(), lr=args.lr, alpha=0.9) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.epochs, eta_min=1e-7, last_epoch=-1) model.to(device) # 72 subjects in total, during CV, training has 70, testing has 2, one epi, one nonepi train_features = features_numpy[train_ind, :, :] train_adj = adj_numpy[train_ind, :, :] train_labels = labels_numpy[train_ind] state = np.random.get_state() np.random.shuffle(train_features) np.random.set_state(state) np.random.shuffle(train_adj) np.random.set_state(state) np.random.shuffle(train_labels) test_features = features[test_ind, :, :] test_adj = adj[test_ind, :, :] test_labels = labels[test_ind] acc_test = [] start_epoch = 13 gap = 1 mode_on = args.mode for epoch in range(args.epochs): train(epoch, model, optimizer, scheduler, torch.from_numpy(train_features).float().to(device), torch.from_numpy(train_adj).float().to(device), torch.from_numpy(train_labels).long().to(device)) if (epoch >= start_epoch) and (epoch % gap == 0) and (mode_on == True): acc_test.append( test(model, test_features, test_adj, test_labels)) ##=================== test_prob, test_accur, test_labels_squeezed = test( model, test_features, test_adj, test_labels) acc_test.append(test_accur) # max_epochs.append(np.argmax(acc_test)*gap + start_epoch) acc.append(np.max(acc_test)) ##============================================= probs_fold.append(test_prob[:, -1]) test_true.append(test_labels_squeezed.cpu().numpy()) # torch.save({'epoch': args.epochs, # 'model_state_dict': model.state_dict(), # 'optimizer_state_dict': optimizer.state_dict(), # }, args.model_path+"model_cv_2002005010050" + str(len(acc)) + ".pth") del model del optimizer # input("any key") # print(acc, max_epochs) # with open('../results/accuracy_0.4thr_15e_5layers.txt', 'w') as f: # f.write(str(acc)) # f.close() probs = np.array(probs_fold).flatten() print(probs) auc = sklearn.metrics.roc_auc_score( np.array(test_true).flatten(), probs) print(auc) # print(np.mean(acc)) acc_iter.append(np.mean(acc)) auc_iter.append(auc) # print(acc_iter, np.mean(acc_iter), np.std(acc_iter)) print("----------Mean AUC-------------") print(auc_iter, np.mean(auc_iter), np.std(auc_iter)) print("Accuracy") print(acc_iter, np.mean(acc_iter), np.std(acc_iter))
def train(**kwargs): """ GCN training --- - the folder you need: - args.path4AffGraph - args.path4node_feat - path4partial_label - these folder would be created: - data/GCN_prediction/label - data/GCN_prediction/logit """ # os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, [0, 1, 2, 3])) t_start = time.time() # 根据命令行参数更新配置 args.parse(**kwargs) # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda:" + str(kwargs["GPU"])) print(device) # 把有改動的參數寫到tensorboard名稱上 if kwargs["debug"] is False: comment_init = '' for k, v in kwargs.items(): comment_init += '|{} '.format(v) writer = SummaryWriter(comment=comment_init) # === set evaluate object for evaluate later IoU = IOUMetric(args.num_class) IoU_CRF = IOUMetric(args.num_class) # === dataset train_dataloader = graph_voc(start_idx=kwargs["start_index"], end_idx=kwargs["end_index"], device=device) # === for each image, do training and testing in the same graph # for ii, (adj_t, features_t, labels_t, rgbxy_t, img_name, label_fg_t, # label_bg_t) in enumerate(train_dataloader): t4epoch = time.time() for ii, data in enumerate(train_dataloader): if data is None: continue # === use RGBXY as feature # if args.use_RGBXY: # data["rgbxy_t"] = normalize_rgbxy(data["rgbxy_t"]) # features_t = data["rgbxy_t"].clone() # === only RGB as feature t_be = time.time() if args.use_lap: """ is constructing................ """ H, W, C = data["rgbxy_t"].shape A = torch.zeros([H * W, H * W], dtype=torch.float64) def find_neibor(card_x, card_y, H, W, radius=2): """ Return idx of neibors of (x,y) in list --- """ neibors_idx = [] for idx_x in np.arange(card_x - radius, card_x + radius + 1): for idx_y in np.arange(card_y - radius, card_y + radius + 1): if (-radius < idx_x < H) and (-radius < idx_y < W): neibors_idx.append( (idx_x * W + idx_y, idx_x, idx_y)) return neibors_idx t_start = time.time() t_start = t4epoch neibors = dict() for node_idx in range(H * W): card_x, card_y = node_idx // W, node_idx % W neibors = find_neibor(card_x, card_y, H, W, radius=1) # print("H:{} W:{} | {} -> ({},{})".format( # H, W, node_idx, card_x, card_y)) for nei in neibors: # print("nei: ", nei) diff_rgb = data["rgbxy_t"][ card_x, card_y, :3] - data["rgbxy_t"][nei[1], nei[2], :3] diff_xy = data["rgbxy_t"][card_x, card_y, 3:] - data["rgbxy_t"][nei[1], nei[2], 3:] A[node_idx, nei[0]] = torch.exp( -torch.pow(torch.norm(diff_rgb), 2) / (2. * args.CRF_deeplab["bi_rgb_std"])) + torch.exp( -torch.pow(torch.norm(diff_xy), 2) / (2. * args.CRF_deeplab["bi_xy_std"])) # print("{:3.1f}s".format(time.time() - t_start)) D = torch.diag(A.sum(dim=1)) L_mat = D - A print("time for Laplacian {:3f} s".format(time.time() - t_be)) # === Model and optimizer img_label = load_image_label_from_xml(img_name=data["img_name"], voc12_root=args.path4VOC_root) img_class = [idx + 1 for idx, f in enumerate(img_label) if int(f) == 1] num_class = np.max(img_class) + 1 # debug("num_class: {} {}".format(num_class + 1, type(num_class + 1)), # line=290) model = GCN( nfeat=data["features_t"].shape[1], nhid=args.num_hid_unit, # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> # image label don't have BG # adaptive num_class should have better performance nclass=args.num_class, # args.num_class| num_class # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> dropout=args.drop_rate) optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # ==== moving tensor to GPU if args.cuda: model.to(device) data["features_t"] = data["features_t"].to(device) data["adj_t"] = data["adj_t"].to(device) data["labels_t"] = data["labels_t"].to(device) data["label_fg_t"] = data["label_fg_t"].to(device) data["label_bg_t"] = data["label_bg_t"].to(device) # L_mat = L_mat.to(device) # === save the prediction before training if args.save_mask_before_train: model.eval() postprocess_image_save(img_name=data["img_name"], model_output=model(data["features_t"], data["adj_t"]).detach(), epoch=0) # ==== Train model # t4epoch = time.time() criterion_ent = HLoss() # criterion_sym = symmetricLoss() for epoch in range(args.max_epoch): model.train() optimizer.zero_grad() output = model(data["features_t"], data["adj_t"]) # === seperate FB/BG label loss_fg = F.nll_loss(output, data["label_fg_t"], ignore_index=255) loss_bg = F.nll_loss(output, data["label_bg_t"], ignore_index=255) # F.log_softmax(label_fg_t, dim=1) # loss_sym = criterion_sym(output, labels_t, ignore_index=255) loss = loss_fg + loss_bg if args.use_ent: loss_entmin = criterion_ent(output, data["labels_t"], ignore_index=255) loss += 10. * loss_entmin if args.use_lap: loss_lap = torch.trace( torch.mm(output.transpose(1, 0), torch.mm(L_mat.type_as(output), output))) / (H * W) gamma = 1e-2 loss += gamma * loss_lap # loss = F.nll_loss(output, labels_t, ignore_index=255) if loss is None: print("skip this image: ", data["img_name"]) break # === for normalize cut # lamda = args.lamda # n_cut = 0. # if args.use_regular_NCut: # W = gaussian_propagator(output) # d = torch.sum(W, dim=1) # for k in range(output.shape[1]): # s = output[idx_test_t, k] # n_cut = n_cut + torch.mm( # torch.mm(torch.unsqueeze(s, 0), W), # torch.unsqueeze(1 - s, 1)) / (torch.dot(d, s)) # === calculus loss & updated parameters # loss_train = loss.cuda() + lamda * n_cut loss_train = loss.cuda() loss_train.backward() optimizer.step() # === save predcit mask at max epoch & IoU of img if (epoch + 1) % args.max_epoch == 0 and args.save_mask: t_now = time.time() if not kwargs["debug"]: evaluate_IoU(model=model, features=data["features_t"], adj=data["adj_t"], img_name=data["img_name"], epoch=args.max_epoch, img_idx=ii + 1, writer=writer, IoU=IoU, IoU_CRF=IoU_CRF, use_CRF=False, save_prediction_np=True) print("[{}/{}] time: {:.4f}s\n\n".format( ii + 1, len(train_dataloader), t_now - t4epoch)) t4epoch = t_now # end for epoch # print( # "loss: {} | loss_fg: {} | loss_bg:{} | loss_entmin: {} | loss_lap: {}" # .format(loss.data.item(), loss_fg.data.item(), loss_bg.data.item(), # loss_entmin.data.item(), loss_lap.data.item())) # end for dataloader if kwargs["debug"] is False: writer.close() print("training was Finished!") print("Total time elapsed: {:.0f} h {:.0f} m {:.0f} s\n".format( (time.time() - t_start) // 3600, (time.time() - t_start) / 60 % 60, (time.time() - t_start) % 60))
def gcn_train(**kwargs): """ GCN training --- - the folder you need: - args.path4AffGraph - args.path4node_feat - path4partial_label - these folder would be created: - data/GCN4DeepLab/Label - data/GCN4DeepLab/Logit """ t_start = time.time() # update config args.parse(**kwargs) device = torch.device("cuda:" + str(kwargs["GPU"])) print(device) # tensorboard if args.use_TB: time_now = datetime.datetime.today() time_now = "{}-{}-{}|{}-{}".format(time_now.year, time_now.month, time_now.day, time_now.hour, time_now.minute // 30) keys_ignore = ["start_index", "GPU"] comment_init = '' for k, v in kwargs.items(): if k not in keys_ignore: comment_init += '|{} '.format(v) writer = SummaryWriter( logdir='runs/{}/{}'.format(time_now, comment_init)) # initial IoUMetric object for evaluation IoU = IOUMetric(args.num_class) # initial dataset train_dataloader = graph_voc(start_idx=kwargs["start_index"], end_idx=kwargs["end_index"], device=device) # train a seperate GCN for each image t4epoch = time.time() for ii, data in enumerate(train_dataloader): if data is None: continue img_label = load_image_label_from_xml(img_name=data["img_name"], voc12_root=args.path4VOC_root) img_class = [idx + 1 for idx, f in enumerate(img_label) if int(f) == 1] num_class = np.max(img_class) + 1 model = GCN(nfeat=data["features_t"].shape[1], nhid=args.num_hid_unit, nclass=args.num_class, dropout=args.drop_rate) optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # put data into GPU if args.cuda: model.to(device) data["features_t"] = data["features_t"].to(device) data["adj_t"] = data["adj_t"].to(device) data["labels_t"] = data["labels_t"].to(device) data["label_fg_t"] = data["label_fg_t"].to(device) data["label_bg_t"] = data["label_bg_t"].to(device) t_be = time.time() H, W, C = data["rgbxy_t"].shape N = H * W # laplacian if args.use_lap: L_mat = compute_lap_test(data, device, radius=2).to(device) print("Time for laplacian {:3.1f} s".format(time.time() - t_be)) criterion_ent = HLoss() for epoch in range(args.max_epoch): model.train() optimizer.zero_grad() output = model(data["features_t"], data["adj_t"]) # foreground and background loss loss_fg = F.nll_loss(output, data["label_fg_t"], ignore_index=255) loss_bg = F.nll_loss(output, data["label_bg_t"], ignore_index=255) loss = loss_fg + loss_bg if args.use_ent: loss_entmin = criterion_ent(output, data["labels_t"], ignore_index=255) loss += 10. * loss_entmin if args.use_lap: loss_lap = torch.trace( torch.mm(output.transpose(1, 0), torch.mm(L_mat.type_as(output), output))) / N gamma = 1e-2 loss += gamma * loss_lap if loss is None: print("skip this image: ", data["img_name"]) break loss_train = loss.cuda() loss_train.backward() optimizer.step() # save predicted mask and IoU at max epoch if (epoch + 1) % args.max_epoch == 0 and args.save_mask: t_now = time.time() evaluate_IoU(model=model, features=data["features_t"], adj=data["adj_t"], img_name=data["img_name"], img_idx=ii + 1, writer=writer, IoU=IoU, save_prediction_np=True) print("evaluate time: {:3.1f} s".format(time.time() - t_now)) print("[{}/{}] time: {:.1f}s\n\n".format( ii + 1, len(train_dataloader), t_now - t4epoch)) t4epoch = t_now print("======================================") if writer is not None: writer.close() print("training was Finished!") print("Total time elapsed: {:.0f} h {:.0f} m {:.0f} s\n".format( (time.time() - t_start) // 3600, (time.time() - t_start) / 60 % 60, (time.time() - t_start) % 60))
def main(args): # convert boolean type for args assert args.use_ist in ['True', 'False'], ["Only True or False for use_ist, get ", args.use_ist] assert args.split_input in ['True', 'False'], ["Only True or False for split_input, get ", args.split_input] assert args.split_output in ['True', 'False'], ["Only True or False for split_output, get ", args.split_output] assert args.self_loop in ['True', 'False'], ["Only True or False for self_loop, get ", args.self_loop] assert args.use_layernorm in ['True', 'False'], ["Only True or False for use_layernorm, get ", args.use_layernorm] assert args.use_random_proj in ['True', 'False'], ["Only True or False for use_random_proj, get ", args.use_random_proj] use_ist = (args.use_ist == 'True') split_input = (args.split_input == 'True') split_output = (args.split_output == 'True') self_loop = (args.self_loop == 'True') use_layernorm = (args.use_layernorm == 'True') use_random_proj = (args.use_random_proj == 'True') # make sure hidden layer is the correct shape assert (args.n_hidden % args.num_subnet) == 0 # load and preprocess dataset global t0 if args.dataset in {'cora', 'citeseer', 'pubmed'}: data = load_data(args) else: raise NotImplementedError(f'{args.dataset} is not a valid dataset') # randomly project the input to make it dense if use_random_proj: # densify input features with random projection from sklearn import random_projection # make sure input features are divisible by number of subnets # otherwise some parameters of the last subnet will be handled improperly n_components = int(data.features.shape[-1] / args.num_subnet) * args.num_subnet transformer = random_projection.GaussianRandomProjection(n_components=n_components) new_feature = transformer.fit_transform(data.features) features = torch.FloatTensor(new_feature) else: assert (data.features.shape[-1] % args.num_subnet) == 0. features = torch.FloatTensor(data.features) labels = torch.LongTensor(data.labels) train_mask = torch.ByteTensor(data.train_mask) val_mask = torch.ByteTensor(data.val_mask) test_mask = torch.ByteTensor(data.test_mask) in_feats = features.shape[1] n_classes = data.num_labels n_edges = data.graph.number_of_edges() print("""----Data statistics------' #Edges %d #Classes %d #Train samples %d #Val samples %d #Test samples %d""" % (n_edges, n_classes, train_mask.sum().item(), val_mask.sum().item(), test_mask.sum().item())) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') features = features.to(device) labels = labels.to(device) train_mask = train_mask.to(device) val_mask = val_mask.to(device) test_mask = test_mask.to(device) # graph preprocess and calculate normalization factor g = data.graph # add self loop if self_loop: g.remove_edges_from(nx.selfloop_edges(g)) g.add_edges_from(zip(g.nodes(), g.nodes())) g = DGLGraph(g) g = g.to(device) n_edges = g.number_of_edges() # normalization degs = g.in_degrees().float() norm = torch.pow(degs, -0.5) norm[torch.isinf(norm)] = 0 norm = norm.to(device) g.ndata['norm'] = norm.unsqueeze(1) # create GCN model model = GCN( g, in_feats, args.n_hidden, n_classes, args.n_layers, F.relu, args.dropout, use_layernorm) model = model.to(device) loss_fcn = torch.nn.CrossEntropyLoss() # initialize graph dur = [] record = [] sub_models = [] opt_list = [] sub_dict_list = [] main_dict = None for epoch in range(args.n_epochs): if epoch >= 3: t0 = time.time() if use_ist: model.eval() # IST training: # Distribute parameter to sub networks num_subnet = args.num_subnet if (epoch % args.iter_per_site) == 0.: main_dict = model.state_dict() feats_idx = [] # store all layer indices within a single list # create input partition if split_input: feats_idx.append(torch.chunk(torch.randperm(in_feats), num_subnet)) else: feats_idx.append(None) # create hidden layer partitions for i in range(1, args.n_layers): feats_idx.append(torch.chunk(torch.randperm(args.n_hidden), num_subnet)) # create output layer partitions if split_output: feats_idx.append(torch.chunk(torch.randperm(args.n_hidden), num_subnet)) else: feats_idx.append(None) for subnet_id in range(args.num_subnet): if (epoch % args.iter_per_site) == 0.: # create the sub model to train sub_model = GCN( g, in_feats, args.n_hidden, n_classes, args.n_layers, F.relu, args.dropout, use_layernorm, split_input, split_output, args.num_subnet) sub_model = sub_model.to(device) sub_dict = main_dict.copy() # split input params if split_input: idx = feats_idx[0][subnet_id] sub_dict['layers.0.weight'] = main_dict['layers.0.weight'][idx, :] # split hidden params (and output params) for i in range(1, args.n_layers + 1): if i == args.n_layers and not split_output: pass # params stay the same else: idx = feats_idx[i][subnet_id] sub_dict[f'layers.{i - 1}.weight'] = sub_dict[f'layers.{i -1}.weight'][:, idx] sub_dict[f'layers.{i - 1}.bias'] = main_dict[f'layers.{i - 1}.bias'][idx] sub_dict[f'layers.{i}.weight'] = main_dict[f'layers.{i}.weight'][idx, :] # use a lr scheduler curr_lr = args.lr if epoch >= int(args.n_epochs*0.5): curr_lr /= 10 if epoch >= int(args.n_epochs*0.75): curr_lr /= 10 # import params into subnet for training sub_model.load_state_dict(sub_dict) sub_models.append(sub_model) sub_models = sub_models[-num_subnet:] optimizer = torch.optim.Adam( sub_model.parameters(), lr=curr_lr, weight_decay=args.weight_decay) opt_list.append(optimizer) opt_list = opt_list[-num_subnet:] else: sub_model = sub_models[subnet_id] optimizer = opt_list[subnet_id] # train a sub network optimizer.zero_grad() sub_model.train() if split_input: model_input = features[:, feats_idx[0][subnet_id]] else: model_input = features logits = sub_model(model_input) loss = loss_fcn(logits[train_mask], labels[train_mask]) # reset optimization for every sub training loss.backward() optimizer.step() # save sub model parameter if ( ((epoch + 1) % args.iter_per_site == 0.) or (epoch == args.n_epochs - 1)): sub_dict = sub_model.state_dict() sub_dict_list.append(sub_dict) sub_dict_list = sub_dict_list[-num_subnet:] # Merge parameter to main network: # force aggregation if training about to end if ( ((epoch + 1) % args.iter_per_site == 0.) or (epoch == args.n_epochs - 1)): #keys = main_dict.keys() update_dict = main_dict.copy() # copy in the input parameters if split_input: if args.n_layers <= 1 and not split_output: for idx, sub_dict in zip(feats_idx[0], sub_dict_list): update_dict['layers.0.weight'][idx, :] = sub_dict['layers.0.weight'] else: for i, sub_dict in enumerate(sub_dict_list): curr_idx = feats_idx[0][i] next_idx = feats_idx[1][i] correct_rows = update_dict['layers.0.weight'][curr_idx, :] correct_rows[:, next_idx] = sub_dict['layers.0.weight'] update_dict['layers.0.weight'][curr_idx, :] = correct_rows else: if args.n_layers <= 1 and not split_output: update_dict['layers.0.weight'] = sum(sub_dict['layers.0.weight'] for sub_dict in sub_dict_list) / len(sub_dict_list) else: for i, sub_dict in enumerate(sub_dict_list): next_idx = feats_idx[1][i] update_dict['layers.0.weight'][:, next_idx] = sub_dict['layers.0.weight'] # copy the rest of the parameters for i in range(1, args.n_layers + 1): if i == args.n_layers: if not split_output: update_dict[f'layers.{i-1}.bias'] = sum(sub_dict[f'layers.{i-1}.bias'] for sub_dict in sub_dict_list) / len(sub_dict_list) update_dict[f'layers.{i}.weight'] = sum(sub_dict[f'layers.{i}.weight'] for sub_dict in sub_dict_list) / len(sub_dict_list) else: for idx, sub_dict in zip(feats_idx[i], sub_dict_list): update_dict[f'layers.{i-1}.bias'][idx] = sub_dict[f'layers.{i-1}.bias'] update_dict[f'layers.{i}.weight'][idx, :] = sub_dict[f'layers.{i}.weight'] else: if i >= args.n_layers - 1 and not split_output: for idx, sub_dict in zip(feats_idx[i], sub_dict_list): update_dict[f'layers.{i-1}.bias'][idx] = sub_dict[f'layers.{i-1}.bias'] update_dict[f'layers.{i}.weight'][idx, :] = sub_dict[f'layers.{i}.weight'] else: for idx, sub_dict in enumerate(sub_dict_list): curr_idx = feats_idx[i][idx] next_idx = feats_idx[i+1][idx] update_dict[f'layers.{i-1}.bias'][curr_idx] = sub_dict[f'layers.{i-1}.bias'] correct_rows = update_dict[f'layers.{i}.weight'][curr_idx, :] correct_rows[:, next_idx] = sub_dict[f'layers.{i}.weight'] update_dict[f'layers.{i}.weight'][curr_idx, :] = correct_rows model.load_state_dict(update_dict) else: raise NotImplementedError('Should train with IST') if epoch >= 3: dur.append(time.time() - t0) acc_val = evaluate(model, features, labels, val_mask) acc_test = evaluate(model, features, labels, test_mask) print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Val Accuracy {:.4f} | Test Accuracy {:.4f} |" "ETputs(KTEPS) {:.2f}".format(epoch, np.mean(dur), loss.item(), acc_val, acc_test, n_edges / np.mean(dur) / 1000)) record.append([acc_val, acc_test]) all_test_acc = [v[1] for v in record] all_val_acc = [v[0] for v in record] acc = evaluate(model, features, labels, test_mask) print(f"Final Test Accuracy: {acc:.4f}") print(f"Best Val Accuracy: {max(all_val_acc):.4f}") print(f"Best Test Accuracy: {max(all_test_acc):.4f}")
gen_config["type"] = gen_type gen_config["neg_ratio"] = args.neg_ratio if gen_type == "lsm": gen_config["hidden_x"] = args.hidden_x if gen_type == "sbm": gen_config["p0"] = args.p0 gen_config["p1"] = args.p1 post_config = copy.deepcopy(model_args) post_config["type"] = post_type if post_type == "gat": post_config["num_heads"] = args.num_heads post_config["hidden_size"] = int(args.hidden / args.num_heads) model = GenGNN(gen_config, post_config) model = model.to(args.device) optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) if hasattr(model, "gen"): def train_loss_fn(model, data): post_y_log_prob = model(data) nll_generative = model.gen.nll_generative(data, post_y_log_prob) nll_discriminative = F.nll_loss(post_y_log_prob[data.train_mask], data.y[data.train_mask]) return nll_generative + args.lamda * nll_discriminative else: def train_loss_fn(model, data):