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):
