def visualize():
adj, features, y_train, y_val, y_test, train_mask, val_mask, test_mask = load_data(
args.dataset)
model = GCN(features.shape[1], args.hidden_dim, y_train.shape[1], 0)
gcn_layer1 = model.gcn_layer1
gcn_layer1.register_forward_hook(hook_fn_forward)
if args.checkpoint:
model.load_state_dict(torch.load(args.checkpoint))
model.eval()
model(adj, features)
x_all = gcn_layer1_output[0]
y_train = np.argmax(y_train[train_mask, :].numpy(), axis=1)
y_val = np.argmax(y_val[val_mask, :].numpy(), axis=1)
y_test = np.argmax(y_test[test_mask, :].numpy(), axis=1)
tsne_vis(x_all[train_mask], y_train, classnames[args.dataset], 'train_set',
args.dataset)
tsne_vis(x_all[val_mask], y_val, classnames[args.dataset], 'val_set',
args.dataset)
tsne_vis(x_all[test_mask], y_test, classnames[args.dataset], 'test_set',
args.dataset)
def main():
global args
args = parser.parse_args()
cfg = load_config(args.config)
device = torch.device("cuda:0" if not args.no_cuda else "cpu")
if args.mode == 'train':
exp_name = 'gcn' + '_old_wc1_small_bs_4_lr_1e-2_test'
tb_writer = SummaryWriter(tb_log_dir + exp_name)
train_loader, val_loader = build_data_loader(args.mode)
model = GCN(cfg).to(device)
optimizer = optim.Adam(model.parameters(),
args.learning_rate, (0.9, 0.999),
eps=1e-08)
#optimizer = optim.SGD(model.parameters(), args.learning_rate, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer,
step_size=50,
gamma=0.5)
total_tb_it = 0
best_val_loss = 1000
for epoch in range(args.epochs):
scheduler.step(epoch)
total_tb_it = train(cfg, model, train_loader, device, optimizer,
epoch, tb_writer, total_tb_it)
val_loss = validate(cfg, model, val_loader, device, tb_writer,
total_tb_it)
if val_loss <= best_val_loss:
best_val_loss = val_loss
name = checkpoint_dir + exp_name + '_model.pkl'
state = {
'epoch': epoch,
'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict()
}
torch.save(state, name)
tb_writer.close()
elif args.mode == 'predict':
print('Load data...')
test_loader = build_data_loader(args.mode)
print('Start predicting...')
model = GCN(cfg).to(device)
model.load_state_dict(torch.load(args.resume_file)['model_state'])
#summary(model, [(200, 1), (200, 200)])
test(cfg, model, test_loader, device)
else:
raise Exception("Unrecognized mode.")
save_path = os.path.join(args.paths["project_path"],
"saves{}".format(args.eval["version"]),
"{}-{}.pth".format(args.model, trained_epochs))
dataset = TestDataset()
if args.model == "GCN":
model = GCN(cityscapes.num_classes, args.img_size, k=args.K).cuda()
elif args.model == "GCN_DENSENET":
model = GCN_DENSENET(cityscapes.num_classes, args.img_size,
k=args.K).cuda()
elif args.model == "GCN_DECONV":
model = GCN_DECONV(cityscapes.num_classes, args.img_size,
k=args.K).cuda()
elif args.model == "GCN_PSP":
model = GCN_PSP(cityscapes.num_classes, args.img_size, k=args.K).cuda()
elif args.model == "GCN_COMB":
model = GCN_COMBINED(cityscapes.num_classes, args.img_size).cuda()
elif args.model == "GCN_RESNEXT":
model = GCN_RESNEXT(dataset.num_classes, k=self.args.K).cuda()
elif args.model == "UNet":
model = UNet(cityscapes.num_classes).cuda()
elif args.model == "PSPNet":
model = PSPNet(cityscapes.num_classes, args.img_size).cuda()
else:
raise ValueError("Invalid model arg.")
model.load_state_dict(torch.load(save_path))
evaluator = Evaluator(dataset)
evaluator.eval(model)
metrics=["miou"],
crf=False)
# setup model:
if args.model == "GCN":
model = GCN(dataset.num_classes, dataset.img_size, k=args.K).cuda()
elif args.model == "UNet":
model = UNet(dataset.num_classes).cuda()
else:
raise ValueError("Invalid model arg.")
model.train()
if args.resume:
assert os.path.exists(os.path.join(paths["project_path"], "saves"))
resume_path = save_path.format(args.resume_epoch)
model.load_state_dict(torch.load(resume_path))
start_epoch = args.resume_epoch
else:
start_epoch = 0
# setup loss and optimizer
if args.optimizer == "SGD":
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=True)
elif args.optimizer == "Adam":
optimizer = optim.Adam(model.parameters(),
lr=args.lr,