momentum=0.9,
weight_decay=args.weight_decay)
class_weight = 45117 / torch.tensor([30160, 2, 9757, 1004, 4, 205, 833, 4, 252, 21, 7, 1366, 1323, 83, 56, 26, 1, 13]).cuda()
criterion = nn.CrossEntropyLoss(weight=class_weight)
for current_epoch in range(model.epochs):
if current_epoch % args.step_update_lr == args.step_update_lr-1:
for g in optimizer.param_groups:
g['lr'] = g['lr']/3
model.epoch = current_epoch
print("Training epoch...")
model.train_epoch(train_loader_T,train_loader, optimizer,criterion)
print("Validating epoch...")
conf_matrix = model.val_epoch(val_loader,criterion)
model.visualize_graph()
if model.val_history["loss"][-1] < model.min_val:
print("Saving model...")
model.min_val = model.val_history["loss"][-1]
torch.save(model.state_dict(), model.session_name+"stage_1.pth")
df = pd.DataFrame(conf_matrix)
filepath = 'val_conf_matrix.xlsx'
## Initializing the model
config = CONFIGS["ViT-B_16"]
num_classes = 18
model = VisionTransformer(config,
args.input_dim,
zero_head=True,
num_classes=num_classes)
model.epochs = args.epochs
model.session_name = args.session_name
model.load_pretrained(model.session_name + "stage_1.pth")
model.cuda()
#
print("testing epoch...")
model.epoch = 0
model.epochs = 0
conf_matrix_test, mIoU_test = model.evaluate(test_loader)
metrics_test = compute_metrics(conf_matrix_test)
#
# metrics_test = np.hstack([metrics_test, mIoU_test[:, np.newaxis]])
#
# df = pd.DataFrame(metrics_test)
# filepath = args.session_name + 'test_metrics.xlsx'
# df.to_excel(filepath, index=False)
#
# df = pd.DataFrame(conf_matrix_test)
# filepath = args.session_name + 'test_conf_matrix.xlsx'
# df.to_excel(filepath, index=False)
#