def main(opts):
# ===== Setup distributed =====
distributed.init_process_group(backend='nccl', init_method='env://')
if opts.device is not None:
device_id = opts.device
else:
device_id = opts.local_rank
device = torch.device(device_id)
rank, world_size = distributed.get_rank(), distributed.get_world_size()
if opts.device is not None:
torch.cuda.set_device(opts.device)
else:
torch.cuda.set_device(device_id)
# ===== Initialize logging =====
logdir_full = f"{opts.logdir}/{opts.dataset}/{opts.name}/"
if rank == 0:
logger = Logger(logdir_full,
rank=rank,
debug=opts.debug,
summary=opts.visualize)
else:
logger = Logger(logdir_full,
rank=rank,
debug=opts.debug,
summary=False)
logger.print(f"Device: {device}")
checkpoint_path = f"checkpoints/{opts.dataset}/{opts.name}.pth"
os.makedirs(f"checkpoints/{opts.dataset}", exist_ok=True)
# ===== Setup random seed to reproducibility =====
torch.manual_seed(opts.random_seed)
torch.cuda.manual_seed(opts.random_seed)
np.random.seed(opts.random_seed)
random.seed(opts.random_seed)
# ===== Set up dataset =====
train_dst, val_dst = get_dataset(opts, train=True)
train_loader = data.DataLoader(train_dst,
batch_size=opts.batch_size,
sampler=DistributedSampler(
train_dst,
num_replicas=world_size,
rank=rank),
num_workers=opts.num_workers,
drop_last=True,
pin_memory=True)
val_loader = data.DataLoader(val_dst,
batch_size=opts.batch_size,
sampler=DistributedSampler(
val_dst,
num_replicas=world_size,
rank=rank),
num_workers=opts.num_workers)
logger.info(f"Dataset: {opts.dataset}, Train set: {len(train_dst)}, "
f"Val set: {len(val_dst)}, n_classes {opts.num_classes}")
logger.info(f"Total batch size is {opts.batch_size * world_size}")
# This is necessary for computing the scheduler decay
opts.max_iter = opts.max_iter = opts.epochs * len(train_loader)
# ===== Set up model and ckpt =====
model = Trainer(device, logger, opts)
model.distribute()
cur_epoch = 0
if opts.continue_ckpt:
opts.ckpt = checkpoint_path
if opts.ckpt is not None:
assert os.path.isfile(
opts.ckpt), "Error, ckpt not found. Check the correct directory"
checkpoint = torch.load(opts.ckpt, map_location="cpu")
cur_epoch = checkpoint["epoch"] + 1
model.load_state_dict(checkpoint["model_state"])
logger.info("[!] Model restored from %s" % opts.ckpt)
del checkpoint
else:
logger.info("[!] Train from scratch")
# ===== Train procedure =====
# print opts before starting training to log all parameters
logger.add_table("Opts", vars(opts))
# uncomment if you want qualitative on val
# if rank == 0 and opts.sample_num > 0:
# sample_ids = np.random.choice(len(val_loader), opts.sample_num, replace=False) # sample idxs for visualization
# logger.info(f"The samples id are {sample_ids}")
# else:
# sample_ids = None
label2color = utils.Label2Color(cmap=utils.color_map(
opts.dataset)) # convert labels to images
denorm = utils.Denormalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225
]) # de-normalization for original images
train_metrics = StreamSegMetrics(opts.num_classes)
val_metrics = StreamSegMetrics(opts.num_classes)
results = {}
# check if random is equal here.
logger.print(torch.randint(0, 100, (1, 1)))
while cur_epoch < opts.epochs and not opts.test:
# ===== Train =====
start = time.time()
epoch_loss = model.train(cur_epoch=cur_epoch,
train_loader=train_loader,
metrics=train_metrics,
print_int=opts.print_interval)
train_score = train_metrics.get_results()
end = time.time()
len_ep = int(end - start)
logger.info(
f"End of Epoch {cur_epoch}/{opts.epochs}, Average Loss={epoch_loss[0] + epoch_loss[1]:.4f}, "
f"Class Loss={epoch_loss[0]:.4f}, Reg Loss={epoch_loss[1]}\n"
f"Train_Acc={train_score['Overall Acc']:.4f}, Train_Iou={train_score['Mean IoU']:.4f} "
f"\n -- time: {len_ep // 60}:{len_ep % 60} -- ")
logger.info(
f"I will finish in {len_ep * (opts.epochs - cur_epoch) // 60} minutes"
)
logger.add_scalar("E-Loss", epoch_loss[0] + epoch_loss[1], cur_epoch)
# logger.add_scalar("E-Loss-reg", epoch_loss[1], cur_epoch)
# logger.add_scalar("E-Loss-cls", epoch_loss[0], cur_epoch)
# ===== Validation =====
if (cur_epoch + 1) % opts.val_interval == 0:
logger.info("validate on val set...")
val_loss, _ = model.validate(loader=val_loader,
metrics=val_metrics,
ret_samples_ids=None)
val_score = val_metrics.get_results()
logger.print("Done validation")
logger.info(
f"End of Validation {cur_epoch}/{opts.epochs}, Validation Loss={val_loss}"
)
log_val(logger, val_metrics, val_score, val_loss, cur_epoch)
# keep the metric to print them at the end of training
results["V-IoU"] = val_score['Class IoU']
results["V-Acc"] = val_score['Class Acc']
# ===== Save Model =====
if rank == 0:
if not opts.debug:
save_ckpt(checkpoint_path, model, cur_epoch)
logger.info("[!] Checkpoint saved.")
cur_epoch += 1
torch.distributed.barrier()
# ==== TESTING =====
logger.info("*** Test the model on all seen classes...")
# make data loader
test_dst = get_dataset(opts, train=False)
test_loader = data.DataLoader(test_dst,
batch_size=opts.batch_size_test,
sampler=DistributedSampler(
test_dst,
num_replicas=world_size,
rank=rank),
num_workers=opts.num_workers)
if rank == 0 and opts.sample_num > 0:
sample_ids = np.random.choice(len(test_loader),
opts.sample_num,
replace=False) # sample idxs for visual.
logger.info(f"The samples id are {sample_ids}")
else:
sample_ids = None
val_loss, ret_samples = model.validate(loader=test_loader,
metrics=val_metrics,
ret_samples_ids=sample_ids)
val_score = val_metrics.get_results()
conf_matrixes = val_metrics.get_conf_matrixes()
logger.print("Done test on all")
logger.info(f"*** End of Test on all, Total Loss={val_loss}")
logger.info(val_metrics.to_str(val_score))
log_samples(logger, ret_samples, denorm, label2color, 0)
logger.add_figure("Test_Confusion_Matrix_Recall",
conf_matrixes['Confusion Matrix'])
logger.add_figure("Test_Confusion_Matrix_Precision",
conf_matrixes["Confusion Matrix Pred"])
results["T-IoU"] = val_score['Class IoU']
results["T-Acc"] = val_score['Class Acc']
results["T-Prec"] = val_score['Class Prec']
logger.add_results(results)
logger.add_scalar("T_Overall_Acc", val_score['Overall Acc'])
logger.add_scalar("T_MeanIoU", val_score['Mean IoU'])
logger.add_scalar("T_MeanAcc", val_score['Mean Acc'])
ret = val_score['Mean IoU']
logger.close()
return ret
def main():
opts = get_argparser().parse_args()
opts = modify_command_options(opts)
# Set up visualization
vis = Visualizer(port=opts.vis_port,
env=opts.vis_env) if opts.enable_vis else None
if vis is not None: # display options
vis.vis_table("Options", vars(opts))
os.environ['CUDA_VISIBLE_DEVICES'] = opts.gpu_id
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device: %s" % device)
# Set up random seed
torch.manual_seed(opts.random_seed)
torch.cuda.manual_seed(opts.random_seed)
np.random.seed(opts.random_seed)
random.seed(opts.random_seed)
# Set up dataloader
train_dst, val_dst = get_dataset(opts)
train_loader = data.DataLoader(train_dst,
batch_size=opts.batch_size,
shuffle=True,
num_workers=opts.num_workers)
val_loader = data.DataLoader(
val_dst,
batch_size=opts.batch_size if opts.crop_val else 1,
shuffle=False,
num_workers=opts.num_workers)
print("Dataset: %s, Train set: %d, Val set: %d" %
(opts.dataset, len(train_dst), len(val_dst)))
# Set up model
print("Backbone: %s" % opts.backbone)
model = DeepLabv3(num_classes=opts.num_classes,
backbone=opts.backbone,
pretrained=True,
momentum=opts.bn_mom,
output_stride=opts.output_stride,
use_separable_conv=opts.use_separable_conv)
if opts.use_gn == True:
print("[!] Replace BatchNorm with GroupNorm!")
model = utils.convert_bn2gn(model)
if opts.fix_bn == True:
model.fix_bn()
if torch.cuda.device_count() > 1: # Parallel
print("%d GPU parallel" % (torch.cuda.device_count()))
model = torch.nn.DataParallel(model)
model_ref = model.module # for ckpt
else:
model_ref = model
model = model.to(device)
# Set up metrics
metrics = StreamSegMetrics(opts.num_classes)
# Set up optimizer
decay_1x, no_decay_1x = model_ref.group_params_1x()
decay_10x, no_decay_10x = model_ref.group_params_10x()
optimizer = torch.optim.SGD(params=[
{
"params": decay_1x,
'lr': opts.lr,
'weight_decay': opts.weight_decay
},
{
"params": no_decay_1x,
'lr': opts.lr
},
{
"params": decay_10x,
'lr': opts.lr * 10,
'weight_decay': opts.weight_decay
},
{
"params": no_decay_10x,
'lr': opts.lr * 10
},
],
lr=opts.lr,
momentum=opts.momentum,
nesterov=not opts.no_nesterov)
del decay_1x, no_decay_1x, decay_10x, no_decay_10x
if opts.lr_policy == 'poly':
scheduler = utils.PolyLR(optimizer,
max_iters=opts.epochs * len(train_loader),
power=opts.lr_power)
elif opts.lr_policy == 'step':
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer,
step_size=opts.lr_decay_step,
gamma=opts.lr_decay_factor)
print("Optimizer:\n%s" % (optimizer))
utils.mkdir('checkpoints')
# Restore
best_score = 0.0
cur_epoch = 0
if opts.ckpt is not None and os.path.isfile(opts.ckpt):
checkpoint = torch.load(opts.ckpt)
model_ref.load_state_dict(checkpoint["model_state"])
optimizer.load_state_dict(checkpoint["optimizer_state"])
scheduler.load_state_dict(checkpoint["scheduler_state"])
cur_epoch = checkpoint["epoch"] + 1
best_score = checkpoint['best_score']
print("Model restored from %s" % opts.ckpt)
del checkpoint # free memory
else:
print("[!] Retrain")
def save_ckpt(path):
""" save current model
"""
state = {
"epoch": cur_epoch,
"model_state": model_ref.state_dict(),
"optimizer_state": optimizer.state_dict(),
"scheduler_state": scheduler.state_dict(),
"best_score": best_score,
}
torch.save(state, path)
print("Model saved as %s" % path)
# Set up criterion
criterion = utils.get_loss(opts.loss_type)
#========== Train Loop ==========#
vis_sample_id = np.random.randint(
0, len(val_loader), opts.vis_sample_num,
np.int32) if opts.enable_vis else None # sample idxs for visualization
label2color = utils.Label2Color(cmap=utils.color_map(
opts.dataset)) # convert labels to images
denorm = utils.Denormalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224,
0.225]) # denormalization for ori images
while cur_epoch < opts.epochs:
# ===== Train =====
model.train()
if opts.fix_bn == True:
model_ref.fix_bn()
epoch_loss = train(cur_epoch=cur_epoch,
criterion=criterion,
model=model,
optim=optimizer,
train_loader=train_loader,
device=device,
scheduler=scheduler,
vis=vis)
print("End of Epoch %d/%d, Average Loss=%f" %
(cur_epoch, opts.epochs, epoch_loss))
if opts.enable_vis:
vis.vis_scalar("Epoch Loss", cur_epoch, epoch_loss)
# ===== Save Latest Model =====
if (cur_epoch + 1) % opts.ckpt_interval == 0:
save_ckpt('checkpoints/latest_%s_%s.pkl' %
(opts.backbone, opts.dataset))
# ===== Validation =====
if (cur_epoch + 1) % opts.val_interval == 0:
print("validate on val set...")
model.eval()
val_score, ret_samples = validate(model=model,
loader=val_loader,
device=device,
metrics=metrics,
ret_samples_ids=vis_sample_id)
print(metrics.to_str(val_score))
# ===== Save Best Model =====
if val_score['Mean IoU'] > best_score: # save best model
best_score = val_score['Mean IoU']
save_ckpt('checkpoints/best_%s_%s.pkl' %
(opts.backbone, opts.dataset))
if vis is not None: # visualize validation score and samples
vis.vis_scalar("[Val] Overall Acc", cur_epoch,
val_score['Overall Acc'])
vis.vis_scalar("[Val] Mean IoU", cur_epoch,
val_score['Mean IoU'])
vis.vis_table("[Val] Class IoU", val_score['Class IoU'])
for k, (img, target, lbl) in enumerate(ret_samples):
img = (denorm(img) * 255).astype(np.uint8)
target = label2color(target).transpose(2, 0,
1).astype(np.uint8)
lbl = label2color(lbl).transpose(2, 0, 1).astype(np.uint8)
concat_img = np.concatenate((img, target, lbl),
axis=2) # concat along width
vis.vis_image('Sample %d' % k, concat_img)
if opts.val_on_trainset == True: # validate on train set
print("validate on train set...")
model.eval()
train_score, _ = validate(model=model,
loader=train_loader,
device=device,
metrics=metrics)
print(metrics.to_str(train_score))
if vis is not None:
vis.vis_scalar("[Train] Overall Acc", cur_epoch,
train_score['Overall Acc'])
vis.vis_scalar("[Train] Mean IoU", cur_epoch,
train_score['Mean IoU'])
cur_epoch += 1
def main():
opts = get_argparser().parse_args()
opts = modify_command_options(opts)
os.environ['CUDA_VISIBLE_DEVICES'] = opts.gpu_id
device = torch.device( 'cuda' if torch.cuda.is_available() else 'cpu' )
print("Device: %s"%device)
# Set up random seed
torch.manual_seed(opts.random_seed)
torch.cuda.manual_seed(opts.random_seed)
np.random.seed(opts.random_seed)
random.seed(opts.random_seed)
# Set up dataloader
_, val_dst = get_dataset(opts)
val_loader = data.DataLoader(val_dst, batch_size=opts.batch_size if opts.crop_val else 1 , shuffle=False, num_workers=opts.num_workers)
print("Dataset: %s, Val set: %d"%(opts.dataset, len(val_dst)))
# Set up model
print("Backbone: %s"%opts.backbone)
model = DeepLabv3(num_classes=opts.num_classes, backbone=opts.backbone, pretrained=True, momentum=opts.bn_mom, output_stride=opts.output_stride, use_separable_conv=opts.use_separable_conv)
if opts.use_gn==True:
print("[!] Replace BatchNorm with GroupNorm!")
model = utils.convert_bn2gn(model)
if torch.cuda.device_count()>1: # Parallel
print("%d GPU parallel"%(torch.cuda.device_count()))
model = torch.nn.DataParallel(model)
model_ref = model.module # for ckpt
else:
model_ref = model
model = model.to(device)
# Set up metrics
metrics = StreamSegMetrics(opts.num_classes)
if opts.save_path is not None:
utils.mkdir(opts.save_path)
# Restore
if opts.ckpt is not None and os.path.isfile(opts.ckpt):
checkpoint = torch.load(opts.ckpt)
model_ref.load_state_dict(checkpoint["model_state"])
print("Model restored from %s"%opts.ckpt)
else:
print("[!] Retrain")
label2color = utils.Label2Color(cmap=utils.color_map(opts.dataset)) # convert labels to images
denorm = utils.Denormalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]) # denormalization for ori images
model.eval()
metrics.reset()
idx = 0
if opts.save_path is not None:
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
with torch.no_grad():
for i, (images, labels) in tqdm( enumerate( val_loader ) ):
images = images.to(device, dtype=torch.float32)
labels = labels.to(device, dtype=torch.long)
outputs = model(images)
preds = outputs.detach().max(dim=1)[1].cpu().numpy()
targets = labels.cpu().numpy()
metrics.update(targets, preds)
if opts.save_path is not None:
for i in range(len(images)):
image = images[i].detach().cpu().numpy()
target = targets[i]
pred = preds[i]
image = (denorm(image) * 255).transpose(1,2,0).astype(np.uint8)
target = label2color(target).astype(np.uint8)
pred = label2color(pred).astype(np.uint8)
Image.fromarray(image).save(os.path.join(opts.save_path, '%d_image.png'%idx) )
Image.fromarray(target).save(os.path.join(opts.save_path, '%d_target.png'%idx) )
Image.fromarray(pred).save(os.path.join(opts.save_path, '%d_pred.png'%idx) )
fig = plt.figure()
plt.imshow(image)
plt.axis('off')
plt.imshow(pred, alpha=0.7)
ax = plt.gca()
ax.xaxis.set_major_locator(matplotlib.ticker.NullLocator())
ax.yaxis.set_major_locator(matplotlib.ticker.NullLocator())
plt.savefig(os.path.join(opts.save_path, '%d_overlay.png'%idx), bbox_inches='tight', pad_inches=0)
plt.close()
idx+=1
score = metrics.get_results()
print(metrics.to_str(score))
if opts.save_path is not None:
with open(os.path.join(opts.save_path, 'score.txt'), mode='w') as f:
f.write(metrics.to_str(score))
def main(opts):
distributed.init_process_group(backend='nccl', init_method='env://')
device_id, device = opts.local_rank, torch.device(opts.local_rank)
rank, world_size = distributed.get_rank(), distributed.get_world_size()
torch.cuda.set_device(device_id)
# Initialize logging
task_name = f"{opts.task}-{opts.dataset}"
logdir_full = f"{opts.logdir}/{task_name}/{opts.name}/"
if rank == 0:
logger = Logger(logdir_full, rank=rank, debug=opts.debug, summary=opts.visualize, step=opts.step)
else:
logger = Logger(logdir_full, rank=rank, debug=opts.debug, summary=False)
logger.print(f"Device: {device}")
# Set up random seed
torch.manual_seed(opts.random_seed)
torch.cuda.manual_seed(opts.random_seed)
np.random.seed(opts.random_seed)
random.seed(opts.random_seed)
# xxx Set up dataloader
train_dst, val_dst, test_dst, n_classes = get_dataset(opts)
# reset the seed, this revert changes in random seed
random.seed(opts.random_seed)
train_loader = data.DataLoader(train_dst, batch_size=opts.batch_size,
sampler=DistributedSampler(train_dst, num_replicas=world_size, rank=rank),
num_workers=opts.num_workers, drop_last=True)
val_loader = data.DataLoader(val_dst, batch_size=opts.batch_size if opts.crop_val else 1,
sampler=DistributedSampler(val_dst, num_replicas=world_size, rank=rank),
num_workers=opts.num_workers)
logger.info(f"Dataset: {opts.dataset}, Train set: {len(train_dst)}, Val set: {len(val_dst)},"
f" Test set: {len(test_dst)}, n_classes {n_classes}")
logger.info(f"Total batch size is {opts.batch_size * world_size}")
# xxx Set up model
logger.info(f"Backbone: {opts.backbone}")
step_checkpoint = None
model = make_model(opts, classes=tasks.get_per_task_classes(opts.dataset, opts.task, opts.step))
logger.info(f"[!] Model made with{'out' if opts.no_pretrained else ''} pre-trained")
if opts.step == 0: # if step 0, we don't need to instance the model_old
model_old = None
else: # instance model_old
model_old = make_model(opts, classes=tasks.get_per_task_classes(opts.dataset, opts.task, opts.step - 1))
if opts.fix_bn:
model.fix_bn()
logger.debug(model)
# xxx Set up optimizer
params = []
if not opts.freeze:
params.append({"params": filter(lambda p: p.requires_grad, model.body.parameters()),
'weight_decay': opts.weight_decay})
params.append({"params": filter(lambda p: p.requires_grad, model.head.parameters()),
'weight_decay': opts.weight_decay})
params.append({"params": filter(lambda p: p.requires_grad, model.cls.parameters()),
'weight_decay': opts.weight_decay})
optimizer = torch.optim.SGD(params, lr=opts.lr, momentum=0.9, nesterov=True)
if opts.lr_policy == 'poly':
scheduler = utils.PolyLR(optimizer, max_iters=opts.epochs * len(train_loader), power=opts.lr_power)
elif opts.lr_policy == 'step':
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=opts.lr_decay_step, gamma=opts.lr_decay_factor)
else:
raise NotImplementedError
logger.debug("Optimizer:\n%s" % optimizer)
if model_old is not None:
[model, model_old], optimizer = amp.initialize([model.to(device), model_old.to(device)], optimizer,
opt_level=opts.opt_level)
model_old = DistributedDataParallel(model_old)
else:
model, optimizer = amp.initialize(model.to(device), optimizer, opt_level=opts.opt_level)
# Put the model on GPU
model = DistributedDataParallel(model, delay_allreduce=True)
# xxx Load old model from old weights if step > 0!
if opts.step > 0:
# get model path
if opts.step_ckpt is not None:
path = opts.step_ckpt
else:
path = f"checkpoints/step/{task_name}_{opts.name}_{opts.step - 1}.pth"
# generate model from path
if os.path.exists(path):
step_checkpoint = torch.load(path, map_location="cpu")
model.load_state_dict(step_checkpoint['model_state'], strict=False) # False because of incr. classifiers
if opts.init_balanced:
# implement the balanced initialization (new cls has weight of background and bias = bias_bkg - log(N+1)
model.module.init_new_classifier(device)
# Load state dict from the model state dict, that contains the old model parameters
model_old.load_state_dict(step_checkpoint['model_state'], strict=True) # Load also here old parameters
logger.info(f"[!] Previous model loaded from {path}")
# clean memory
del step_checkpoint['model_state']
elif opts.debug:
logger.info(f"[!] WARNING: Unable to find of step {opts.step - 1}! Do you really want to do from scratch?")
else:
raise FileNotFoundError(path)
# put the old model into distributed memory and freeze it
for par in model_old.parameters():
par.requires_grad = False
model_old.eval()
# xxx Set up Trainer
trainer_state = None
# if not first step, then instance trainer from step_checkpoint
if opts.step > 0 and step_checkpoint is not None:
if 'trainer_state' in step_checkpoint:
trainer_state = step_checkpoint['trainer_state']
# instance trainer (model must have already the previous step weights)
trainer = Trainer(model, model_old, device=device, opts=opts, trainer_state=trainer_state,
classes=tasks.get_per_task_classes(opts.dataset, opts.task, opts.step))
# xxx Handle checkpoint for current model (model old will always be as previous step or None)
best_score = 0.0
cur_epoch = 0
if opts.ckpt is not None and os.path.isfile(opts.ckpt):
checkpoint = torch.load(opts.ckpt, map_location="cpu")
model.load_state_dict(checkpoint["model_state"], strict=True)
optimizer.load_state_dict(checkpoint["optimizer_state"])
scheduler.load_state_dict(checkpoint["scheduler_state"])
cur_epoch = checkpoint["epoch"] + 1
best_score = checkpoint['best_score']
logger.info("[!] Model restored from %s" % opts.ckpt)
# if we want to resume training, resume trainer from checkpoint
if 'trainer_state' in checkpoint:
trainer.load_state_dict(checkpoint['trainer_state'])
del checkpoint
else:
if opts.step == 0:
logger.info("[!] Train from scratch")
# xxx Train procedure
# print opts before starting training to log all parameters
logger.add_table("Opts", vars(opts))
if rank == 0 and opts.sample_num > 0:
sample_ids = np.random.choice(len(val_loader), opts.sample_num, replace=False) # sample idxs for visualization
logger.info(f"The samples id are {sample_ids}")
else:
sample_ids = None
label2color = utils.Label2Color(cmap=utils.color_map(opts.dataset)) # convert labels to images
denorm = utils.Denormalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]) # de-normalization for original images
TRAIN = not opts.test
val_metrics = StreamSegMetrics(n_classes)
results = {}
# check if random is equal here.
logger.print(torch.randint(0,100, (1,1)))
# train/val here
while cur_epoch < opts.epochs and TRAIN:
# ===== Train =====
model.train()
epoch_loss = trainer.train(cur_epoch=cur_epoch, optim=optimizer,
train_loader=train_loader, scheduler=scheduler, logger=logger)
logger.info(f"End of Epoch {cur_epoch}/{opts.epochs}, Average Loss={epoch_loss[0]+epoch_loss[1]},"
f" Class Loss={epoch_loss[0]}, Reg Loss={epoch_loss[1]}")
# ===== Log metrics on Tensorboard =====
logger.add_scalar("E-Loss", epoch_loss[0]+epoch_loss[1], cur_epoch)
logger.add_scalar("E-Loss-reg", epoch_loss[1], cur_epoch)
logger.add_scalar("E-Loss-cls", epoch_loss[0], cur_epoch)
# ===== Validation =====
if (cur_epoch + 1) % opts.val_interval == 0:
logger.info("validate on val set...")
model.eval()
val_loss, val_score, ret_samples = trainer.validate(loader=val_loader, metrics=val_metrics,
ret_samples_ids=sample_ids, logger=logger)
logger.print("Done validation")
logger.info(f"End of Validation {cur_epoch}/{opts.epochs}, Validation Loss={val_loss[0]+val_loss[1]},"
f" Class Loss={val_loss[0]}, Reg Loss={val_loss[1]}")
logger.info(val_metrics.to_str(val_score))
# ===== Save Best Model =====
if rank == 0: # save best model at the last iteration
score = val_score['Mean IoU']
# best model to build incremental steps
save_ckpt(f"checkpoints/step/{task_name}_{opts.name}_{opts.step}.pth",
model, trainer, optimizer, scheduler, cur_epoch, score)
logger.info("[!] Checkpoint saved.")
# ===== Log metrics on Tensorboard =====
# visualize validation score and samples
logger.add_scalar("V-Loss", val_loss[0]+val_loss[1], cur_epoch)
logger.add_scalar("V-Loss-reg", val_loss[1], cur_epoch)
logger.add_scalar("V-Loss-cls", val_loss[0], cur_epoch)
logger.add_scalar("Val_Overall_Acc", val_score['Overall Acc'], cur_epoch)
logger.add_scalar("Val_MeanIoU", val_score['Mean IoU'], cur_epoch)
logger.add_table("Val_Class_IoU", val_score['Class IoU'], cur_epoch)
logger.add_table("Val_Acc_IoU", val_score['Class Acc'], cur_epoch)
# logger.add_figure("Val_Confusion_Matrix", val_score['Confusion Matrix'], cur_epoch)
# keep the metric to print them at the end of training
results["V-IoU"] = val_score['Class IoU']
results["V-Acc"] = val_score['Class Acc']
for k, (img, target, lbl) in enumerate(ret_samples):
img = (denorm(img) * 255).astype(np.uint8)
target = label2color(target).transpose(2, 0, 1).astype(np.uint8)
lbl = label2color(lbl).transpose(2, 0, 1).astype(np.uint8)
concat_img = np.concatenate((img, target, lbl), axis=2) # concat along width
logger.add_image(f'Sample_{k}', concat_img, cur_epoch)
cur_epoch += 1
# ===== Save Best Model at the end of training =====
if rank == 0 and TRAIN: # save best model at the last iteration
# best model to build incremental steps
save_ckpt(f"checkpoints/step/{task_name}_{opts.name}_{opts.step}.pth",
model, trainer, optimizer, scheduler, cur_epoch, best_score)
logger.info("[!] Checkpoint saved.")
torch.distributed.barrier()
# xxx From here starts the test code
logger.info("*** Test the model on all seen classes...")
# make data loader
test_loader = data.DataLoader(test_dst, batch_size=opts.batch_size if opts.crop_val else 1,
sampler=DistributedSampler(test_dst, num_replicas=world_size, rank=rank),
num_workers=opts.num_workers)
# load best model
if TRAIN:
model = make_model(opts, classes=tasks.get_per_task_classes(opts.dataset, opts.task, opts.step))
# Put the model on GPU
model = DistributedDataParallel(model.cuda(device))
ckpt = f"checkpoints/step/{task_name}_{opts.name}_{opts.step}.pth"
checkpoint = torch.load(ckpt, map_location="cpu")
model.load_state_dict(checkpoint["model_state"])
logger.info(f"*** Model restored from {ckpt}")
del checkpoint
trainer = Trainer(model, None, device=device, opts=opts)
model.eval()
val_loss, val_score, _ = trainer.validate(loader=test_loader, metrics=val_metrics, logger=logger)
logger.print("Done test")
logger.info(f"*** End of Test, Total Loss={val_loss[0]+val_loss[1]},"
f" Class Loss={val_loss[0]}, Reg Loss={val_loss[1]}")
logger.info(val_metrics.to_str(val_score))
logger.add_table("Test_Class_IoU", val_score['Class IoU'])
logger.add_table("Test_Class_Acc", val_score['Class Acc'])
logger.add_figure("Test_Confusion_Matrix", val_score['Confusion Matrix'])
results["T-IoU"] = val_score['Class IoU']
results["T-Acc"] = val_score['Class Acc']
logger.add_results(results)
logger.add_scalar("T_Overall_Acc", val_score['Overall Acc'], opts.step)
logger.add_scalar("T_MeanIoU", val_score['Mean IoU'], opts.step)
logger.add_scalar("T_MeanAcc", val_score['Mean Acc'], opts.step)
logger.close()
def test_color_map():
cm = u.color_map()
print(cm)