def get_transform_seg(dataset_name, enc, augment, height, width):
dataMeta = get_dataset_metadata(dataset_name)
return _SegmentationTransform(enc=enc,
augment=augment,
height=height,
width=width,
relabel_target_val=dataMeta.IGNORE_INDEX)
def get_class_weights(dataset_name,
data_list,
enc,
height,
width,
num_workers,
batch_size,
video_train_kth_frame=None):
dataset_meta = get_dataset_metadata(dataset_name)
transform = get_transform_labels_only(dataset_name, enc, height, width)
dataset = get_dataset_labels_only(dataset_name, data_list, transform,
'train', video_train_kth_frame)
loader = DataLoader(dataset,
num_workers=num_workers,
batch_size=batch_size,
shuffle=False)
hist = torch.zeros(dataset_meta.NUM_CLASSES, dtype=torch.float32)
for labels in loader:
hist = hist + torch.histc(labels.float(), dataset_meta.NUM_CLASSES, 0,
dataset_meta.NUM_CLASSES - 1)
hist = hist / hist.sum()
weights = torch.log1p(hist + 0.1).reciprocal()
if dataset_meta.IGNORE_INDEX:
weights[dataset_meta.IGNORE_INDEX] = 0.0
return weights
def trainer(args, seg_network, flow_network, gru_cell):
datasetMeta = get_dataset_metadata(args.dataset)
weight = get_class_weights(args.dataset, args.datalist, False,
TRAIN_RESOLUTION[args.dataset].height,
TRAIN_RESOLUTION[args.dataset].width,
args.num_workers, args.batch_size,
args.kth_frame).to(device)
print(f'Weights are: {weight}')
criterion = nn.CrossEntropyLoss(ignore_index=datasetMeta.IGNORE_INDEX,
reduction='mean',
weight=weight)
gru_optimizer = Adam(params=gru_cell.parameters(),
lr=1e-5,
betas=(0.95, 0.99),
eps=1e-8)
seg_network.requires_grad_(False)
seg_params = []
for module in seg_network.named_modules():
if isinstance(module[1], non_bottleneck_1d):
seg_params += list(module[1].parameters())
module[1].requires_grad_(True)
seg_optimizer = Adam(params=seg_params, lr=1e-10)
ctx = TrainContext(seg_network=seg_network,
flow_network=flow_network,
gru_network=gru_cell,
seg_optimizer=seg_optimizer,
gru_optimizer=gru_optimizer)
return train(args, ctx, criterion)
def segmentation_to_image(seg):
"""seg should be np.array of form height x width x nr_classes"""
assert seg.shape[2] < seg.shape[0] < seg.shape[1]
if seg.shape[2] == 15:
dataset_meta = get_dataset_metadata('v_kitti')
elif seg.shape[2] == 20:
dataset_meta = get_dataset_metadata('cityscapes')
else:
dataset_meta = None
seg = seg.argmax(axis=2)
res = np.zeros((*seg.shape, 3), dtype=np.uint8) # rgb image
for cls in dataset_meta.CLASSES:
res[seg == cls.TRAIN_ID] = cls.COLOR
return res
def main(args):
validate_args(args)
if not args.no_benchmark:
cudnn.benchmark = True
print(f'cudnn.enabled={cudnn.enabled}, cudnn.benchmark={cudnn.benchmark}')
start_training = time.time()
savedir = args.savedir
if not exists(savedir):
makedirs(savedir)
with open(savedir + '/opts.txt', "w") as myfile:
myfile.write(str(args).replace(', ', '\n').replace('(', '(\n'))
copyfile(__file__, savedir + '/' + basename(__file__))
datasetMeta = get_dataset_metadata(args.dataset)
if args.static == 'erfnet':
seg_network = ERFNet(datasetMeta.NUM_CLASSES)
copy_object_sourcefile(seg_network, savedir)
seg_network = nn.DataParallel(seg_network).cuda() # need this for now
status = seg_network.load_state_dict(torch.load(args.static_weights),
strict=False)
print(f'segmentation model loading: {status}')
seg_network = seg_network.module
seg_network.to(device)
if args.flow == 'vcn':
flow_network = load_vcn(args.flow_weights,
VCN_MAXDISP_FAC[args.dataset].maxdisp,
VCN_MAXDISP_FAC[args.dataset].fac,
TRAIN_RESOLUTION[args.dataset].width,
TRAIN_RESOLUTION[args.dataset].height)
copy_object_sourcefile(flow_network, savedir)
flow_network.to(device)
if not args.flow:
flow_network = nn.Module().to(device)
gru_network = GRU(conv_size=(args.gru_size, args.gru_size),
nr_channels=datasetMeta.NUM_CLASSES,
show_timeout=args.show_timeout)
copy_object_sourcefile(gru_network, savedir)
gru_network.to(device)
trainer(args, seg_network, flow_network, gru_network)
training_duration = time.time() - start_training
minutes, seconds = divmod(int(training_duration), 60)
hours, minutes = divmod(minutes, 60)
print(f"Training duration: {hours:02}:{minutes:02}:{seconds:02}")
def trainer(args, model, enc, epoch_callback_before=None):
datasetMeta = get_dataset_metadata(args.dataset)
res = TRAIN_RESOLUTION[args.dataset + '_video']
weight = get_class_weights(args.dataset, args.datalist, enc, res.height, res.width, args.num_workers, args.batch_size, args.kth_frame).to(device)
print(f'Weights are: {weight}')
criterion = nn.CrossEntropyLoss(ignore_index=datasetMeta.IGNORE_INDEX, reduction='mean', weight=weight)
optimizer = Adam(model.parameters(), lr=5e-4, weight_decay=1e-4, betas=(0.9, 0.999))
lam = lambda epoch: pow((1 - (epoch / args.num_epochs)), 0.9)
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lam)
return train(args, model, enc, criterion, optimizer, scheduler, epoch_callback_before)
def train(args, model, enc, criterion, optimizer, scheduler, epoch_callback_before=None):
max_iou = 0
loader_train, loader_val = get_DataLoaders(args, enc)
print(f'Dataset size train: {len(loader_train.dataset)}, val: {len(loader_val.dataset)}')
savedir = args.savedir
logger = AutomatedLogger(savedir, enc)
start_epoch = 1
if args.resume:
max_iou, start_epoch = load_train_resume(enc, model, optimizer, savedir, scheduler)
datasetMeta = get_dataset_metadata(args.dataset)
for epoch in range(start_epoch, args.num_epochs + 1):
if epoch_callback_before:
epoch_callback_before(epoch, model)
curr_lr = 0
for param_group in optimizer.param_groups:
curr_lr = float(param_group['lr'])
print("CURRENT LR: ", param_group['lr'])
iouTrain = IoUMeter(datasetMeta.NUM_CLASSES, datasetMeta.IGNORE_INDEX) if args.train_iou else None
average_epoch_loss_train = do_train_epoch(args, model, optimizer, scheduler, criterion, loader_train, enc, epoch, iouTrain)
iouVal = IoUMeter(datasetMeta.NUM_CLASSES, datasetMeta.IGNORE_INDEX) if args.val_iou else None
average_epoch_loss_val = do_val_epoch(args, model, criterion, loader_val, enc, epoch, iouVal)
logger.write(epoch, average_epoch_loss_train, average_epoch_loss_val, iouTrain.getIoU()[0] if iouTrain else 0, iouVal.getIoU()[0] if iouVal else 0, curr_lr, 0)
if iouVal is None:
curr_iou = -average_epoch_loss_val
else:
curr_iou = iouVal.getIoU()[0]
is_best = curr_iou > max_iou
max_iou = max(curr_iou, max_iou)
save_train_state(args, max_iou, enc, epoch, is_best, model, optimizer, savedir, scheduler)
if is_best:
msg = f'max val iou={iouVal.getIoU()[0] * 100 if iouVal else 0:.2f}, in epoch {epoch}'
if (not enc):
with open(savedir + "/best.txt", "w") as f:
f.write(msg)
else:
with open(savedir + "/best_encoder.txt", "w") as f:
f.write(msg)
return model
def evaluate(args, seg_network, flow_network, gru_cell):
ctx = TrainContext(seg_network=seg_network,
flow_network=flow_network,
gru_network=gru_cell,
seg_optimizer=None,
gru_optimizer=None)
ctx.load_net_state_dicts(args.weights_dir, args.weights_name)
dataset_meta = get_dataset_metadata(args.dataset)
iou = IoUMeter(dataset_meta.NUM_CLASSES,
dataset_meta.IGNORE_INDEX,
resize_first_batch=args.orig_res)
loader = get_DataLoader(args)
print(f'Dataset size train: {len(loader.dataset)}')
print(f'Subset: {args.subset}')
loader = tqdm(loader)
dummy_criterion = lambda _, t: torch.tensor(
0.0, dtype=torch.float32, device=t.device)
train_gru.do_epoch(args,
ctx,
loader,
criterion=dummy_criterion,
train_gru=False,
train_seg=False,
train_flow=False,
IoU=iou)
iou_mean, iou_class = iou.getIoU()
print("IoU per class:")
for iou_class_val, class_name in zip(iou_class, dataset_meta.CLASS_NAMES):
print(f'{iou_class_val * 100:.2f}% {class_name}')
print("=======================================")
print(f'IoU mean: {iou_mean * 100:.2f}%')
if not args.no_out_file:
out_file = join(
args.weights_dir,
f'acc_{args.subset}_{"orig_" if args.orig_res else ""}{iou_mean * 100:.2f}.txt'
)
with open(out_file, 'w') as f:
for i in range(iou_class.shape[0]):
f.write(
f'{iou_class[i] * 100:.2f} {dataset_meta.CLASS_NAMES[i]}\n'
)
f.write(f'{iou_mean * 100:.2f} MEAN')
def main(args):
dataset_meta = get_dataset_metadata(args.dataset)
weights_path = args.weights
assert exists(weights_path), f'{weights_path} does not exist'
model = ERFNet(dataset_meta.NUM_CLASSES)
model = torch.nn.DataParallel(model).to(device)
status = model.load_state_dict(torch.load(weights_path), strict=False)
print(status)
model.eval()
torch.set_grad_enabled(False)
res = TRAIN_RESOLUTION[args.dataset + '_video']
transform = GenericTransformSegEval(
resolution_input=res,
resolution_target=None if args.orig_res else res,
relabel_target_val=dataset_meta.IGNORE_INDEX)
dataset = get_dataset(args.dataset, args.datalist, transform, args.subset)
loader = DataLoader(dataset,
num_workers=args.num_workers,
batch_size=args.batch_size,
shuffle=False)
loader = tqdm(loader)
iou = IoUMeter(dataset_meta.NUM_CLASSES,
dataset_meta.IGNORE_INDEX,
resize_first_batch=args.orig_res)
for step, (images, labels) in enumerate(loader):
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
outputs = outputs.max(1)[1].unsqueeze(1)
iou.addBatch(outputs, labels)
iou_mean, iou_class = iou.getIoU()
print("IoU per class:")
for iou_class_val, class_name in zip(iou_class, dataset_meta.CLASS_NAMES):
print(f'{iou_class_val * 100:.2f}% {class_name}')
print("=======================================")
print(f'IoU mean: {iou_mean * 100:.2f}%')
def get_DataLoader(args):
res = TRAIN_RESOLUTION[args.dataset]
datasetMeta = get_dataset_metadata(args.dataset)
transform = GenericVideoTransform(
resolution_input=res,
resolution_target=None if args.orig_res else res,
relabel_target_val=datasetMeta.IGNORE_INDEX)
dataset = get_dataset(args, transform, args.subset)
pin_mem = not args.no_pin_memory
print(f'Pin memory: {pin_mem}')
loader = DataLoader(dataset,
num_workers=args.num_workers,
batch_size=None,
shuffle=False,
pin_memory=pin_mem)
return loader
def main(args):
if not args.no_benchmark:
cudnn.benchmark = True
print(f'cudnn.enabled={cudnn.enabled}, cudnn.benchmark={cudnn.benchmark}')
start_training = time.time()
savedir = args.savedir
if not exists(savedir):
makedirs(savedir)
with open(savedir + '/opts.txt', "w") as f:
f.write(str(args))
datasetMeta = get_dataset_metadata(args.dataset)
model = ERFNet(datasetMeta.NUM_CLASSES)
copy_object_sourcefile(model, savedir)
copyfile(__file__, savedir + '/' + basename(__file__))
model = torch.nn.DataParallel(model).to(device)
if args.weights:
status = model.load_state_dict(torch.load(args.weights), strict=False)
print(status)
if (not args.decoder):
print("-------TRAINING ENC-------")
model = trainer(args, model, enc=True)
print("-------TRAINING DEC-------")
if (not args.weights):
pretrainedEnc = next(model.children()).encoder
model = ERFNet(datasetMeta.NUM_CLASSES, encoder=pretrainedEnc)
model = torch.nn.DataParallel(model).to(device)
trainer(args, model, enc=False)
training_duration = time.time() - start_training
minutes, seconds = divmod(int(training_duration), 60)
hours, minutes = divmod(minutes, 60)
print(f"Training duration: {hours:02}:{minutes:02}:{seconds:02}")
def main(args):
validate_args(args)
if not args.no_benchmark:
cudnn.benchmark = True
print(f'cudnn.enabled={cudnn.enabled}, cudnn.benchmark={cudnn.benchmark}')
start_eval = time.time()
datasetMeta = get_dataset_metadata(args.dataset)
if args.static == 'erfnet':
seg_network = ERFNet(datasetMeta.NUM_CLASSES)
seg_network.to(device)
if args.flow == 'vcn':
res = TRAIN_RESOLUTION[args.dataset]
vcn_param = VCN_MAXDISP_FAC[args.dataset]
flow_network = VCN_Wrapped(
[1, res.width, res.height],
md=[int(4 * (vcn_param.maxdisp / 256)), 4, 4, 4, 4],
fac=vcn_param.fac,
meanL=None,
meanR=None)
flow_network.to(device)
if not args.flow:
flow_network = nn.Module().to(device)
gru_network = GRU(conv_size=(args.gru_size, args.gru_size),
nr_channels=datasetMeta.NUM_CLASSES,
show_timeout=args.show_timeout)
gru_network.to(device)
evaluate(args, seg_network, flow_network, gru_network)
training_duration = time.time() - start_eval
minutes, seconds = divmod(int(training_duration), 60)
hours, minutes = divmod(minutes, 60)
print(f"Evaluation duration: {hours:02}:{minutes:02}:{seconds:02}")
def get_transform_labels_only(dataset_name, enc, height, width):
dataMeta = get_dataset_metadata(dataset_name)
return _LabelTransform(enc, height, width, dataMeta.IGNORE_INDEX)
def get_transform_video(dataset_name, height, width):
dataMeta = get_dataset_metadata(dataset_name)
res = Resolution(height, width)
return GenericVideoTransform(resolution_input=res,
resolution_target=res,
relabel_target_val=dataMeta.IGNORE_INDEX)
def train(args, ctx: TrainContext, criterion):
ctx.max_iou = 0
loader_train, loader_val = get_data_loaders(args)
print(
f'Dataset size train: {len(loader_train.dataset)}, val: {len(loader_val.dataset)}'
)
logger = AutomatedLogger(args.savedir)
if args.resume:
ctx.load_checkpoint(args.resume)
start_epoch = ctx.last_completed_epoch + 1
datasetMeta = get_dataset_metadata(args.dataset)
for ctx.epoch in range(start_epoch, args.num_epochs + 1):
curr_lr = 0
for param_group in ctx.gru_optimizer.param_groups:
curr_lr = float(param_group['lr'])
print("CURRENT GRU LR: ", param_group['lr'])
train_seg = bool(args.epoch_start_seg
is not None) and ctx.epoch >= args.epoch_start_seg
loader_train.dataset.shuffle()
time_start_epoch = time.time()
iouTrain = IoUMeter(
datasetMeta.NUM_CLASSES,
datasetMeta.IGNORE_INDEX) if args.train_iou else None
average_epoch_loss_train = do_epoch(args,
ctx,
loader_train,
criterion,
train_gru=True,
train_seg=train_seg,
train_flow=False,
IoU=iouTrain)
iouVal = IoUMeter(datasetMeta.NUM_CLASSES,
datasetMeta.IGNORE_INDEX) if args.val_iou else None
average_epoch_loss_val = do_epoch(args,
ctx,
loader_val,
criterion,
train_gru=False,
train_seg=False,
train_flow=False,
IoU=iouVal)
time_epoch = time.time() - time_start_epoch
print(f'TRAIN + VAL duration: {time_epoch}')
logger.write(ctx.epoch, average_epoch_loss_train,
average_epoch_loss_val,
iouTrain.getIoU()[0] if iouTrain else 0,
iouVal.getIoU()[0] if iouVal else 0, curr_lr, time_epoch)
if iouVal is None:
curr_iou = -average_epoch_loss_val
else:
curr_iou = iouVal.getIoU()[0]
is_best = curr_iou > ctx.max_iou
ctx.max_iou = max(curr_iou, ctx.max_iou)
ctx.last_completed_epoch = ctx.epoch
save_train_state(args, ctx, is_best)
if is_best:
with open(args.savedir + "/best.txt", "w") as f:
f.write(
f'max val iou={iouVal.getIoU()[0] * 100 if iouVal else 0:.2f}, in epoch {ctx.epoch}'
)