def train(trainloader,
prob_model,
seg_model,
criterion,
device,
optimizer,
writer_idx,
writer=None):
"""Create the model and start the training."""
# For logging the training status
losses = AverageMeter()
prob_model.train()
seg_model.eval()
with tqdm(total=len(trainloader)) as pbar:
for i_iter, batch in enumerate(tqdm(trainloader)):
# feature = batch['feature'].to(device)
# label = batch['label'].to(device)
images = batch["rgb"].to(device)
masks = batch["mask"].to(device)
optimizer.zero_grad()
# Output probability
# output = model(feature, label)
output_dict = get_output(seg_model, images)
seg_output = output_dict['output']
feature = output_dict['feature']
# Generate mask from the segmentation result
seg_output_amax = torch.argmax(seg_output, dim=1, keepdim=True)
# Forward the feature and calculate probability
prob_output = prob_model(feature)
# Loss calculation
# loss = criterion(prob_output, labels, seg_output_amax==PLANT)
loss = criterion(prob_output, masks)
# Model regularizer
losses.update(loss.item(), feature.size(0))
# Optimise
loss.backward()
optimizer.step()
# Write summary
write_image_to_writer(trainloader,
seg_model,
prob_model,
1.0,
writer,
writer_idx,
mode='train',
device=device)
writer.add_scalar('traversability_mask/train/loss', losses.avg, writer_idx)
writer.add_scalar('traversability_mask/train/learning_rate',
optimizer.param_groups[0]['lr'], writer_idx)
def evaluate(args, model, image_list, seg_classes, device):
im_size = tuple(args.im_size)
# get color map for pascal dataset
if args.dataset == 'pascal':
from utilities.color_map import VOCColormap
cmap = VOCColormap().get_color_map_voc()
else:
cmap = None
model.eval()
inter_meter = AverageMeter()
union_meter = AverageMeter()
miou_class = MIOU(num_classes=seg_classes)
for i, imgName in tqdm(enumerate(image_list)):
img = Image.open(imgName).convert('RGB')
w, h = img.size
img = data_transform(img, im_size)
img = img.unsqueeze(0) # add a batch dimension
img = img.to(device)
img_out = model(img)
img_out = img_out.squeeze(0) # remove the batch dimension
img_out = img_out.max(0)[1].byte() # get the label map
img_out = img_out.to(device='cpu').numpy()
if args.dataset == 'city':
# cityscape uses different IDs for training and testing
# so, change from Train IDs to actual IDs
img_out = relabel(img_out)
img_out = Image.fromarray(img_out)
# resize to original size
img_out = img_out.resize((w, h), Image.NEAREST)
# pascal dataset accepts colored segmentations
if args.dataset == 'pascal':
img_out.putpalette(cmap)
# save the segmentation mask
name = imgName.split('/')[-1]
img_extn = imgName.split('.')[-1]
name = '{}/{}'.format(args.savedir, name.replace(img_extn, 'png'))
img_out.save(name)
def test(testloader, prob_model, seg_model, criterion, device, writer_idx,
writer):
"""Create the model and start the evaluation process."""
# For logging the training status
sigmoid = nn.Sigmoid()
prob_sum_meter = AverageMeter()
losses = AverageMeter()
## model for evaluation
prob_model.eval()
# TODO: Change this (implement the same function in 'utility/utils.py', or uncomment the code below with slight modification)
with torch.no_grad():
# Calculate a constant c
for i_iter, batch in enumerate(tqdm(testloader)):
images = batch["rgb"].to(device)
masks = batch["mask"].to(device)
if args.use_depth:
depths = batch["depth"].to(device)
output_dict = get_output(seg_model, images)
feature = output_dict['feature']
masks = torch.reshape(
masks, (masks.size(0), -1, masks.size(1), masks.size(2)))
prob_output = prob_model(feature)
prob_sum_meter.update(
sigmoid(prob_output)[masks == 1].mean().item(), images.size(0))
# Loss calculation
# loss = criterion(prob_output, labels, seg_output_amax==PLANT)
loss = criterion(prob_output, masks)
# Model regularizer
losses.update(loss.item(), feature.size(0))
c = prob_sum_meter.avg
# Write summary
writer.add_scalar('traversability_mask/test/loss', losses.avg, writer_idx)
write_image_to_writer(testloader,
seg_model,
prob_model,
c,
writer,
writer_idx,
mode='test',
device=device)
return {"c": c, "loss": losses.avg}
def val_seg_per_image(model,
dataset_loader,
criterion=None,
num_classes=21,
device='cuda',
use_depth=False):
model.eval()
inter_meter = AverageMeter()
union_meter = AverageMeter()
batch_time = AverageMeter()
end = time.time()
miou_class = MIOU(num_classes=num_classes - 1)
if criterion:
losses = AverageMeter()
accuracy_list = {}
with torch.no_grad():
for i, batch in enumerate(dataset_loader):
inputs = batch[0].to(device=device)
target = batch[1].to(device=device)
if use_depth:
depth = batch[2].to(device=device)
outputs = model(inputs, depth)
else:
outputs = model(inputs)
if criterion:
if device == 'cuda':
loss = criterion(outputs, target).mean()
if isinstance(outputs, (list, tuple)):
target_dev = outputs[0].device
outputs = gather(outputs, target_device=target_dev)
else:
loss = criterion(outputs, target)
losses.update(loss.item(), inputs.size(0))
inter, union = miou_class.get_iou(outputs, target)
inter_meter.update(inter)
union_meter.update(union)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
loss_ = losses.avg if criterion is not None else 0
# print_log_message("[%d/%d]\t\tBatch Time:%.4f\t\tLoss:%.4f\t\tmiou:%.4f" %
# (i, len(dataset_loader), batch_time.avg, loss_, miou))
accuracy_list[i] = miou
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
print_info_message('Mean IoU: {0:.2f}'.format(miou))
return accuracy_list
def main(args):
crop_size = args.crop_size
assert isinstance(crop_size, tuple)
print_info_message('Running Model at image resolution {}x{} with batch size {}'.format(crop_size[0], crop_size[1],
args.batch_size))
if not os.path.isdir(args.savedir):
os.makedirs(args.savedir)
writer = SummaryWriter(log_dir=args.savedir)
num_gpus = torch.cuda.device_count()
device = 'cuda' if num_gpus > 0 else 'cpu'
from data_loader.segmentation.greenhouse import color_encoding as color_encoding_greenhouse
from data_loader.segmentation.greenhouse import color_palette
from data_loader.segmentation.camvid import color_encoding as color_encoding_camvid
# Outsource
os_model_name_list = [args.os_model1, args.os_model2, args.os_model3]
os_weights_name_list = [args.os_weights1, args.os_weights2, args.os_weights3]
os_data_name_list = [args.outsource1, args.outsource2, args.outsource3]
os_model_name_list = [x for x in os_model_name_list if x is not None]
os_weights_name_list = [x for x in os_weights_name_list if x is not None]
os_data_name_list = [x for x in os_data_name_list if x is not None]
os_model_list = []
print(os_model_name_list)
print(os_weights_name_list)
print(os_data_name_list)
for os_m, os_w, os_d in zip(os_model_name_list, os_weights_name_list, os_data_name_list):
if os_d == 'camvid':
os_seg_classes = 13
elif os_d == 'cityscapes':
os_seg_classes = 20
elif os_d == 'forest' or os_d == 'greenhouse':
os_seg_classes = 5
os_model = import_os_model(args, os_model=os_m, os_weights=os_w, os_seg_classes=os_seg_classes)
os_model_list.append(os_model)
from data_loader.segmentation.greenhouse import GreenhouseRGBDSegmentation, GREENHOUSE_CLASS_LIST
seg_classes = len(GREENHOUSE_CLASS_LIST)
val_dataset = GreenhouseRGBDSegmentation(root='./vision_datasets/greenhouse/', list_name=args.val_list, use_traversable=False,
train=False, size=crop_size, use_depth=args.use_depth,
normalize=args.normalize)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False,
pin_memory=True, num_workers=args.workers)
start_epoch = 0
best_miou = 0.0
losses = AverageMeter()
ce_losses = AverageMeter()
nid_losses = AverageMeter()
batch_time = AverageMeter()
inter_meter = AverageMeter()
union_meter = AverageMeter()
miou_class = MIOU(num_classes=seg_classes)
with torch.no_grad():
for i, batch in enumerate(val_loader):
inputs = batch[0].to(device=device)
target = batch[1].to(device=device)
name = batch[2]
output_list = []
for m, os_data in zip(os_model_list, os_data_name_list):
# Output: Numpy, KLD: Numpy
output, _ = get_output(m, inputs)
output = output.transpose(1,2,0)
amax_output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
# save visualized seg maps & predication prob map
if os_data == 'camvid':
amax_output = id_camvid_to_greenhouse[amax_output]
elif os_data == 'cityscapes':
amax_output = id_cityscapes_to_greenhouse[amax_output]
elif os_data == 'forest':
amax_output = id_forest_to_greenhouse[amax_output]
output_list.append(amax_output)
amax_output = merge_outputs(np.array(output_list),
seg_classes=5, thresh='all')
# Output the generated label images
if args.output_image:
for path_name in name:
# path_name = name[0]
image_name = path_name.split('/')[-1]
image_name = image_name.rsplit('.', 1)[0]
amax_output_img_color = colorize_mask(amax_output, color_palette)
amax_output_img_color.save('%s/%s_color.png' % (args.savedir, image_name))
for output_i, name_i in zip(output_list, os_data_name_list):
amax_output_img_color = colorize_mask(output_i, color_palette)
amax_output_img_color.save('%s/%s_color_%s.png' % (args.savedir, image_name, name_i))
outputs_argmax = torch.from_numpy(amax_output)
inter, union = miou_class.get_iou(outputs_argmax, target)
inter_meter.update(inter)
union_meter.update(union)
# measure elapsed time
print("Batch {}/{} finished".format(i+1, len(val_loader)))
iou = inter_meter.sum / (union_meter.sum + 1e-10) * 100
miou = iou[[1, 2, 3]].mean()
writer.add_scalar('label_eval/IoU', miou, 0)
writer.add_scalar('label_eval/plant', iou[1], 0)
writer.add_scalar('label_eval/artificial_object', iou[2], 0)
writer.add_scalar('label_eval/ground', iou[3], 0)
writer.close()
def test(testloader,
model,
criterion,
device,
optimizer,
class_encoding,
writer_idx,
class_weights=None,
writer=None):
"""Create the model and start the evaluation process."""
## scorer
h, w = map(int, args.input_size.split(','))
test_image_size = (h, w)
# For logging the training status
losses = AverageMeter()
batch_time = AverageMeter()
inter_meter = AverageMeter()
union_meter = AverageMeter()
miou_class = MIOU(num_classes=4)
kld_layer = PixelwiseKLD()
from data_loader.segmentation.greenhouse import GreenhouseRGBDSegmentation
ds = GreenhouseRGBDSegmentation(list_name=args.data_test_list,
train=False,
use_traversable=True,
use_depth=args.use_depth)
testloader = data.DataLoader(ds,
batch_size=32,
shuffle=False,
pin_memory=args.pin_memory)
## model for evaluation
model.eval()
## upsampling layer
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=test_image_size,
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=test_image_size, mode='bilinear')
## output of deeplab is logits, not probability
softmax2d = nn.Softmax2d()
## evaluation process
start_eval = time.time()
total_loss = 0
ious = 0
# TODO: Change this (implement the same function in 'utility/utils.py', or uncomment the code below with slight modification)
# total_loss, (iou, miou) = util.run_validation(model, testloader, criterion, metric, device, writer, interp)
with torch.no_grad():
for index, batch in enumerate(testloader):
#image, label, depth, name, reg_weights = batch
images = batch[0].to(device)
labels = batch[1].to(device)
if args.use_depth:
depths = batch[2].to(device)
reg_weights = batch[4]
else:
reg_weights = batch[3]
# Upsample the output of the model to the input size
# TODO: Change the input according to the model type
interp = None
if interp is not None:
if args.use_depth:
if args.model == 'espdnet':
pred = interp(model(images, depths))
elif args.model == 'espdnetue':
(pred, pred_aux) = interp(model(images, depths))
else:
if args.model == 'espdnet':
pred = interp(model(images))
elif args.model == 'espdnetue':
(pred, pred_aux) = interp(model(images))
elif args.model == 'deeplabv3':
output = model(images)
pred = interp(output['out'])
pred_aux = interp(output['aux'])
else:
if args.use_depth:
if args.model == 'espdnet':
pred = model(images, depths)
elif args.model == 'espdnetue':
(pred, pred_aux) = model(images, depths)
else:
if args.model == 'espdnet':
pred = model(images)
elif args.model == 'espdnetue':
(pred, pred_aux) = model(images)
elif args.model == 'deeplabv3':
output = model(images)
pred = output['out']
pred_aux = output['aux']
loss = criterion(
pred,
labels) # torch.max returns a tuple of (maxvalues, indices)
inter, union = miou_class.get_iou(pred, labels)
inter_meter.update(inter)
union_meter.update(union)
losses.update(loss.item(), images.size(0))
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
writer.add_scalar('traversability_mask/test/mean_IoU', miou, writer_idx)
writer.add_scalar('traversability_mask/test/loss', losses.avg, writer_idx)
writer.add_scalar('traversability_mask/test/traversable_plant_IoU', iou[0],
writer_idx)
writer.add_scalar('traversability_mask/test/other_plant_mean_IoU', iou[1],
writer_idx)
writer.add_scalar('traversability_mask/test/artificial_object_mean_IoU',
iou[2], writer_idx)
writer.add_scalar('traversability_mask/test/ground_mean_IoU', iou[3],
writer_idx)
# if args.dataset == 'greenhouse':
# # TODO: Check
if args.use_depth:
# model, images, depths=None, labels=None, predictions=None, class_encoding=None, writer=None, epoch=None, data=None, device=None
in_training_visualization_img(model,
images=images,
depths=depths,
labels=labels,
class_encoding=class_encoding,
writer=writer,
epoch=writer_idx,
data='traversability_mask/test',
device=device)
else:
in_training_visualization_img(model,
images=images,
labels=labels,
class_encoding=class_encoding,
writer=writer,
epoch=writer_idx,
data='traversability_mask/test',
device=device)
return miou
pred = pred * (target > 0)
# pred = pred * (target < self.num_classes)
inter = pred * (pred == target)
# inter = pred * (target < self.num_classes)
area_inter = torch.histc(inter.float(),
bins=self.num_classes,
min=1,
max=self.num_classes)
area_pred = torch.histc(pred.float(),
bins=self.num_classes,
min=1,
max=self.num_classes)
area_mask = torch.histc(target.float(),
bins=self.num_classes,
min=1,
max=self.num_classes)
area_union = area_pred + area_mask - area_inter + self.epsilon
return area_inter.numpy(), area_union.numpy()
if __name__ == '__main__':
from utilities.utils import AverageMeter
inter = AverageMeter()
union = AverageMeter()
a = torch.Tensor(1, 21, 224, 224).random_(254, 256)
b = torch.Tensor(1, 21, 224, 224).random_(254, 256)
m = MIOU()
print(m.get_iou(a, b))
def train_seg(model,
dataset_loader,
optimizer,
criterion,
num_classes,
epoch,
device='cuda'):
losses = AverageMeter()
batch_time = AverageMeter()
inter_meter = AverageMeter()
union_meter = AverageMeter()
end = time.time()
model.train()
miou_class = MIOU(num_classes=num_classes)
for i, (inputs, target) in enumerate(dataset_loader):
inputs = inputs.to(device=device)
target = target.to(device=device)
outputs = model(inputs)
if device == 'cuda':
loss = criterion(outputs, target).mean()
if isinstance(outputs, (list, tuple)):
target_dev = outputs[0].device
outputs = gather(outputs, target_device=target_dev)
else:
loss = criterion(outputs, target)
inter, union = miou_class.get_iou(outputs, target)
inter_meter.update(inter)
union_meter.update(union)
losses.update(loss.item(), inputs.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0: # print after every 100 batches
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
print_log_message(
"Epoch: %d[%d/%d]\t\tBatch Time:%.4f\t\tLoss:%.4f\t\tmiou:%.4f"
% (epoch, i, len(dataset_loader), batch_time.avg, losses.avg,
miou))
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
return miou, losses.avg
def train(data_loader, model, criterion, optimizer, device, epoch=-1):
model.train()
train_loss = AverageMeter()
train_cl_loss = AverageMeter()
train_loc_loss = AverageMeter()
batch_time = AverageMeter()
end = time.time()
for batch_idx, (images, boxes, labels) in enumerate(data_loader):
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
optimizer.zero_grad()
confidences, locations = model(images)
regression_loss, classification_loss = criterion(
confidences, locations, labels, boxes)
loss = regression_loss + classification_loss
loss.backward()
optimizer.step()
N = images.size(0)
train_loss.update(loss.item(), N)
train_cl_loss.update(classification_loss.item(), N)
train_loc_loss.update(regression_loss.item(), N)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx and batch_idx % 10 == 0:
print_log_message(
"Epoch: %d[%d/%d]\t\tBatch Time:%.4f\t\tCls Loss: %.4f(%.4f)\t \tLoc Loss: %.4f(%.4f)\t\tTotal Loss: %.4f(%.4f)"
% (epoch, batch_idx, len(data_loader), batch_time.avg,
train_cl_loss.val, train_cl_loss.avg, train_loc_loss.val,
train_loc_loss.avg, train_loss.val, train_loss.avg))
return train_loss.avg, train_cl_loss.avg, train_loc_loss.avg
def validate(data_loader, model, criteria=None, device='cuda'):
batch_time = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
if criteria:
losses = AverageMeter()
# switch to evaluate mode
model.eval()
# with torch.no_grad():
end = time.time()
with torch.no_grad():
for i, (input, target) in enumerate(data_loader):
input = input.to(device)
target = target.to(device)
# compute output
output = model(input)
if criteria:
loss = criteria(output, target)
# measure accuracy and record loss
prec1, prec5 = accuracy(output, target, topk=(1, 5))
if criteria:
losses.update(loss.item(), input.size(0))
top1.update(prec1[0].item(), input.size(0))
top5.update(prec5[0].item(), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0 and criteria: # print after every 100 batches
print_log_message(
"Batch:[%d/%d]\t\tBatchTime:%.3f\t\tLoss:%.3f\t\ttop1:%.3f (%.3f)\t\ttop5:%.3f(%.3f)"
% (i, len(data_loader), batch_time.avg, losses.avg,
top1.val, top1.avg, top5.val, top5.avg))
elif i % 10:
print_log_message(
"Batch:[%d/%d]\t\tBatchTime:%.3f\t\ttop1:%.3f (%.3f)\t\ttop5:%.3f(%.3f)"
% (i, len(data_loader), batch_time.avg, top1.val, top1.avg,
top5.val, top5.avg))
print_info_message(' * Prec@1:%.3f Prec@5:%.3f' % (top1.avg, top5.avg))
if criteria:
return top1.avg, losses.avg
else:
return top1.avg
def validate_multi(data_loader, model, criteria, device='cuda'):
batch_time = AverageMeter()
losses = AverageMeter()
prec = AverageMeter()
rec = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
tp, fp, fn, tn, count = 0, 0, 0, 0, 0
tp_size, fn_size = 0, 0
with torch.no_grad():
for i, (input, target) in enumerate(data_loader):
input = input.to(device=device)
target = target.to(device=device)
original_target = target
target = target.max(dim=1)[0]
# compute output
output = model(input)
loss = criteria(output, target.float())
# measure accuracy and record loss
pred = output.data.gt(0.0).long()
tp += (pred + target).eq(2).sum(dim=0)
fp += (pred - target).eq(1).sum(dim=0)
fn += (pred - target).eq(-1).sum(dim=0)
tn += (pred + target).eq(0).sum(dim=0)
three_pred = pred.unsqueeze(1).expand(-1, 3, -1) # n, 3, 80
tp_size += (three_pred + original_target).eq(2).sum(dim=0)
fn_size += (three_pred - original_target).eq(-1).sum(dim=0)
count += input.size(0)
this_tp = (pred + target).eq(2).sum()
this_fp = (pred - target).eq(1).sum()
this_fn = (pred - target).eq(-1).sum()
this_tn = (pred + target).eq(0).sum()
this_acc = (this_tp + this_tn).float() / (
this_tp + this_tn + this_fp + this_fn).float()
this_prec = this_tp.float() / (this_tp + this_fp).float(
) * 100.0 if this_tp + this_fp != 0 else 0.0
this_rec = this_tp.float() / (this_tp + this_fn).float(
) * 100.0 if this_tp + this_fn != 0 else 0.0
losses.update(float(loss), input.size(0))
prec.update(float(this_prec), input.size(0))
rec.update(float(this_rec), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
p_c = [
float(tp[i].float() / (tp[i] + fp[i]).float()) *
100.0 if tp[i] > 0 else 0.0 for i in range(len(tp))
]
r_c = [
float(tp[i].float() / (tp[i] + fn[i]).float()) *
100.0 if tp[i] > 0 else 0.0 for i in range(len(tp))
]
f_c = [
2 * p_c[i] * r_c[i] / (p_c[i] + r_c[i]) if tp[i] > 0 else 0.0
for i in range(len(tp))
]
mean_p_c = sum(p_c) / len(p_c)
mean_r_c = sum(r_c) / len(r_c)
mean_f_c = sum(f_c) / len(f_c)
p_o = tp.sum().float() / (tp + fp).sum().float() * 100.0
r_o = tp.sum().float() / (tp + fn).sum().float() * 100.0
f_o = 2 * p_o * r_o / (p_o + r_o)
if i % 100 == 0:
print_log_message(
'Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Precision {prec.val:.2f} ({prec.avg:.2f})\t'
'Recall {rec.val:.2f} ({rec.avg:.2f})'.format(
i,
len(data_loader),
batch_time=batch_time,
loss=losses,
prec=prec,
rec=rec))
print(
'P_C {:.2f} R_C {:.2f} F_C {:.2f} P_O {:.2f} R_O {:.2f} F_O {:.2f}'
.format(mean_p_c, mean_r_c, mean_f_c, p_o, r_o, f_o))
print_info_message(
' * P_C {:.2f} R_C {:.2f} F_C {:.2f} P_O {:.2f} R_O {:.2f} F_O {:.2f}'
.format(mean_p_c, mean_r_c, mean_f_c, p_o, r_o, f_o))
return f_o, losses.avg
def train(data_loader, model, criteria, optimizer, epoch, device='cuda'):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to train mode
model.train()
end = time.time()
for i, (input, target) in enumerate(data_loader):
# measure data loading time
data_time.update(time.time() - end)
input = input.to(device)
target = target.to(device)
# compute output
output = model(input)
# compute loss
loss = criteria(output, target)
# measure accuracy and record loss
prec1, prec5 = accuracy(output, target, topk=(1, 5))
#losses.update(loss.data[0], input.size(0))
losses.update(loss.item(), input.size(0))
top1.update(prec1[0].item(), input.size(0))
top5.update(prec5[0].item(), input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0: #print after every 100 batches
print_log_message(
"Epoch: %d[%d/%d]\t\tBatch Time:%.4f\t\tLoss:%.4f\t\ttop1:%.4f (%.4f)\t\ttop5:%.4f (%.4f)"
% (epoch, i, len(data_loader), batch_time.avg, losses.avg,
top1.val, top1.avg, top5.val, top5.avg))
return top1.avg, losses.avg
def train_multi(data_loader, model, criteria, optimizer, epoch, device='cuda'):
batch_time = AverageMeter()
losses = AverageMeter()
prec = AverageMeter()
rec = AverageMeter()
# switch to train mode
model.train()
end = time.time()
tp, fp, fn, tn, count = 0, 0, 0, 0, 0
p_o, r_o, f_o = 0.0, 0.0, 0.0
for i, (input, target) in enumerate(data_loader):
target = target.to(device=device)
target = target.max(dim=1)[0]
# compute output
output = model(input)
loss = criteria(output, target.float()) * 80.0
# measure accuracy and record loss
pred = output.gt(0.0).long()
tp += (pred + target).eq(2).sum(dim=0)
fp += (pred - target).eq(1).sum(dim=0)
fn += (pred - target).eq(-1).sum(dim=0)
tn += (pred + target).eq(0).sum(dim=0)
count += input.size(0)
this_tp = (pred + target).eq(2).sum()
this_fp = (pred - target).eq(1).sum()
this_fn = (pred - target).eq(-1).sum()
this_tn = (pred + target).eq(0).sum()
this_acc = (this_tp + this_tn).float() / (this_tp + this_tn + this_fp +
this_fn).float()
this_prec = this_tp.float() / (this_tp + this_fp).float(
) * 100.0 if this_tp + this_fp != 0 else 0.0
this_rec = this_tp.float() / (this_tp + this_fn).float(
) * 100.0 if this_tp + this_fn != 0 else 0.0
losses.update(float(loss), input.size(0))
prec.update(float(this_prec), input.size(0))
rec.update(float(this_rec), input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
p_c = [
float(tp[i].float() / (tp[i] + fp[i]).float()) *
100.0 if tp[i] > 0 else 0.0 for i in range(len(tp))
]
r_c = [
float(tp[i].float() / (tp[i] + fn[i]).float()) *
100.0 if tp[i] > 0 else 0.0 for i in range(len(tp))
]
f_c = [
2 * p_c[i] * r_c[i] / (p_c[i] + r_c[i]) if tp[i] > 0 else 0.0
for i in range(len(tp))
]
p_o = tp.sum().float() / (tp + fp).sum().float() * 100.0
r_o = tp.sum().float() / (tp + fn).sum().float() * 100.0
f_o = 2 * p_o * r_o / (p_o + r_o)
if i % 100 == 0:
print_log_message(
'Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Precision {prec.val:.2f} ({prec.avg:.2f})\t'
'Recall {rec.val:.2f} ({rec.avg:.2f})'.format(
epoch,
i,
len(data_loader),
batch_time=batch_time,
loss=losses,
prec=prec,
rec=rec))
return f_o, losses.avg
def calculate_iou(dataloader,
seg_model,
prob_model,
c,
writer,
thresh=0.5,
writer_idx=None,
device='cuda',
visualize=True):
# For logging the training status
sigmoid = nn.Sigmoid()
iou_sum_meter = AverageMeter()
acc_sum_meter = AverageMeter()
pre_sum_meter = AverageMeter()
rec_sum_meter = AverageMeter()
## model for evaluation
seg_model.eval()
prob_model.eval()
# thresh = c / 2.0
with torch.no_grad():
# Calculate a constant c
# For each data batch
for i_iter, batch in enumerate(tqdm(dataloader)):
images = batch["rgb"].to(device)
masks = batch["mask"].to(device)
masks = torch.reshape(
masks, (masks.size(0), -1, masks.size(1), masks.size(2)))
if args.use_depth:
depths = batch["depth"].to(device)
output_dict = get_output(seg_model, images)
feature = output_dict['feature']
prob_output = prob_model(feature)
prob_output = sigmoid(prob_output) / c
# prob_output /= prob_output.max()
prob_output[prob_output >= 1] = 1.0
# Calculate IoUs with different thresholds
pred_mask = prob_output > thresh
if i_iter == 0:
pred_mask_visualize = torch.zeros(pred_mask.size())
pred_mask_visualize[pred_mask] = 1
union = pred_mask | (masks == 1)
TP = (pred_mask & (masks == 1)).sum().item()
FP = (pred_mask).sum().item() - TP
FN = (masks == 1).sum().item() - TP
TN = (~union).sum().item()
print(TP, FP, FN, TN)
iou = TP / (TP + FP + FN)
acc = (TP + TN) / (TP + TN + FP + FN)
pre = TP / (TP + FP)
rec = TP / (TP + FN)
print(iou, acc, pre, rec)
# print(pred_mask.sum(), (masks == 1).sum())
# print(union.size(), int.size())
# print(union.sum().item(), int.sum().item(), iou.item())
iou_sum_meter.update(iou, feature.size(0))
acc_sum_meter.update(acc, feature.size(0))
pre_sum_meter.update(pre, feature.size(0))
rec_sum_meter.update(rec, feature.size(0))
# Output the best IoU if histogram is not required
writer.add_scalar('traversability_mask/test/IoU', iou_sum_meter.avg,
writer_idx)
writer.add_scalar('traversability_mask/test/Accuracy',
acc_sum_meter.avg, writer_idx)
writer.add_scalar('traversability_mask/test/Precision',
pre_sum_meter.avg, writer_idx)
writer.add_scalar('traversability_mask/test/Recall', rec_sum_meter.avg,
writer_idx)
# Visualize the mask that achieves the best IoU
if visualize:
mask_grid = torchvision.utils.make_grid(
pred_mask_visualize.data.cpu()).numpy()
writer.add_image('traversability_mask/test/pred_mask', mask_grid,
writer_idx)
def calculate_iou_with_different_threshold(dataloader,
seg_model,
prob_model,
c,
writer,
writer_idx=None,
device='cuda',
min_thresh=0.0,
max_thresh=1.0,
step=0.1,
histogram=True,
visualize=True):
# For logging the training status
sigmoid = nn.Sigmoid()
num = round((max_thresh - min_thresh) / step)
iou_sum_meter_list = []
for i in range(num):
iou_sum_meter = AverageMeter()
iou_sum_meter_list.append(iou_sum_meter)
pred_mask_list = [None] * num
## model for evaluation
seg_model.eval()
prob_model.eval()
with torch.no_grad():
# Calculate a constant c
# For each data batch
for i_iter, batch in enumerate(tqdm(dataloader)):
images = batch["rgb"].to(device)
masks = batch["mask"].to(device)
masks = torch.reshape(
masks, (masks.size(0), -1, masks.size(1), masks.size(2)))
if args.use_depth:
depths = batch["depth"].to(device)
output_dict = get_output(seg_model, images)
feature = output_dict['feature']
prob_output = prob_model(feature)
prob_output = sigmoid(prob_output) / c
prob_output /= prob_output.max()
# Calculate IoUs with different thresholds
for i, thresh in enumerate(np.arange(min_thresh, max_thresh,
step)):
pred_mask = prob_output > thresh
if i_iter == 0:
pred_mask_visualize = torch.zeros(pred_mask.size())
pred_mask_visualize[pred_mask] = 1
pred_mask_list[i] = pred_mask_visualize
union = pred_mask | (masks == 1)
inter = pred_mask & (masks == 1)
TP = ((gt == 1) & (pred == 1)).sum().item()
FN = ((gt == 1)).sum().item() - TP
FP = ((pred == 1)).sum().item() - TP
acc = (TP + TN) / torch.numel(gt)
pre = TP / (TP + FP)
rec = TP / (TP + FN)
iou = torch.div(inter.sum().float(), union.sum().float())
# print(pred_mask.sum(), (masks == 1).sum())
# print(union.size(), int.size())
# print(union.sum().item(), int.sum().item(), iou.item())
iou_sum_meter_list[i].update(iou.item(), feature.size(0))
max_iou = 0.0
best_thresh = 0.0
best_index = -1
# Calculate the best IoU and output value in writer if histogram is required
for i, thresh in enumerate(np.arange(min_thresh, max_thresh, step)):
if iou_sum_meter_list[i].avg > max_iou:
max_iou = iou_sum_meter_list[i].avg
best_index = i
if histogram:
writer.add_scalar('traversability_mask/test/iou_per_thresh',
iou_sum_meter_list[i].avg, i)
# Output the best IoU if histogram is not required
if not histogram:
writer.add_scalar('traversability_mask/test/best_iou', max_iou,
writer_idx)
# Visualize the mask that achieves the best IoU
if visualize:
mask_grid = torchvision.utils.make_grid(
pred_mask_list[best_index].data.cpu()).numpy()
writer.add_image('traversability_mask/test/pred_mask', mask_grid,
writer_idx)
def train_seg_cls(model, dataset_loader, optimizer, criterion_seg, num_classes, epoch, criterion_cls, cls_loss_weight=1.0, device='cuda', use_depth=False):
losses = AverageMeter()
cls_losses = AverageMeter()
seg_losses = AverageMeter()
batch_time = AverageMeter()
inter_meter = AverageMeter()
union_meter = AverageMeter()
end = time.time()
model.train()
miou_class = MIOU(num_classes=num_classes)
for i, batch in enumerate(dataset_loader):
inputs = batch[0].to(device=device)
target = batch[1].to(device=device)
if use_depth:
depth = batch[2].to(device=device)
outputs_seg, outputs_cls = model(inputs, depth)
else:
outputs_seg, outputs_cls = model(inputs)
cls_ids = batch[3].to(device=device)
if device == 'cuda':
loss_seg = criterion_seg(outputs_seg, target).mean()
loss_cls = criterion_cls(outputs_cls, cls_ids).mean()
loss = loss_seg + cls_loss_weight * loss_cls
if isinstance(outputs_seg, (list, tuple)):
target_dev = outputs[0].device
outputs_seg = gather(outputs_seg, target_device=target_dev)
else:
loss_seg = criterion_seg(outputs_seg, target)
loss_cls = criterion_cls(outputs_cls, cls_ids)
loss = loss_seg + cls_loss_weight * loss_cls
inter, union = miou_class.get_iou(outputs_seg, target)
inter_meter.update(inter)
union_meter.update(union)
losses.update(loss.item(), inputs.size(0))
seg_losses.update(loss_seg.item(), inputs.size(0))
cls_losses.update(loss_cls.item(), inputs.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0: # print after every 100 batches
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
print_log_message("Epoch: %d[%d/%d]\t\tBatch Time:%.4f\t\tLoss:%.4f\t\tmiou:%.4f" %
(epoch, i, len(dataset_loader), batch_time.avg, losses.avg, miou))
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
return miou, losses.avg, seg_losses.avg, cls_losses.avg
def val_seg_cls(model, dataset_loader, criterion_seg=None, criterion_cls=None, num_classes=21, cls_loss_weight=1.0, device='cuda', use_depth=False):
model.eval()
inter_meter = AverageMeter()
union_meter = AverageMeter()
batch_time = AverageMeter()
end = time.time()
miou_class = MIOU(num_classes=num_classes)
if criterion_seg:
losses = AverageMeter()
cls_losses = AverageMeter()
seg_losses = AverageMeter()
with torch.no_grad():
for i, batch in enumerate(dataset_loader):
inputs = batch[0].to(device=device)
target = batch[1].to(device=device)
if use_depth:
depth = batch[2].to(device=device)
outputs_seg, outputs_cls = model(inputs, depth)
else:
outputs_seg, outputs_cls = model(inputs)
cls_ids = batch[3].to(device=device)
if criterion_seg and criterion_cls:
if device == 'cuda':
loss_seg = criterion_seg(outputs_seg, target).mean()
loss_cls = criterion_cls(outputs_cls, cls_ids).mean()
loss = loss_seg + cls_loss_weight * loss_cls
if isinstance(outputs_seg, (list, tuple)):
target_dev = outputs[0].device
outputs_seg = gather(outputs_seg, target_device=target_dev)
else:
loss_seg = criterion_seg(outputs_seg, target)
loss_cls = criterion_cls(outputs_cls, cls_ids)
loss = loss_seg + cls_loss_weight * loss_cls
losses.update(loss.item(), inputs.size(0))
seg_losses.update(loss_seg.item(), inputs.size(0))
cls_losses.update(loss_cls.item(), inputs.size(0))
inter, union = miou_class.get_iou(outputs_seg, target)
inter_meter.update(inter)
union_meter.update(union)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0: # print after every 100 batches
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
loss_ = losses.avg if criterion_seg is not None else 0
print_log_message("[%d/%d]\t\tBatch Time:%.4f\t\tLoss:%.4f\t\tmiou:%.4f" %
(i, len(dataset_loader), batch_time.avg, loss_, miou))
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
print_info_message('Mean IoU: {0:.2f}'.format(miou))
if criterion_seg and criterion_cls:
return miou, losses.avg, seg_losses.avg, cls_losses.avg
else:
return miou, 0, 0, 0
def train_seg_ue(model,
dataset_loader,
optimizer,
criterion,
num_classes,
epoch,
device='cuda',
use_depth=False,
add_criterion=None,
weight=1.0,
greenhouse_use_trav=False):
losses = AverageMeter()
ce_losses = AverageMeter()
nid_losses = AverageMeter()
batch_time = AverageMeter()
inter_meter = AverageMeter()
union_meter = AverageMeter()
end = time.time()
model.train()
miou_class = MIOU(num_classes=num_classes - 1)
kld_layer = PixelwiseKLD()
print("train_seg_ue()")
b = 0.015
for i, batch in enumerate(dataset_loader):
inputs = batch[0].to(device=device)
target = batch[1].to(device=device)
if use_depth:
depth = batch[2].to(device=device)
outputs = model(inputs, depth)
else:
outputs = model(inputs)
if isinstance(outputs, OrderedDict):
out_aux = outputs['aux']
outputs = outputs['out']
else:
out_aux = outputs[1]
outputs = outputs[0]
kld = kld_layer(outputs, out_aux)
outputs = outputs + 0.5 * out_aux
if device == 'cuda':
# print("Target size {}".format(target.size()))
#
loss = criterion(outputs, target).mean() # + kld.mean()
if add_criterion is not None:
loss2 = add_criterion(inputs, outputs.to(device)) * weight
loss += loss2
if isinstance(outputs, (list, tuple)):
target_dev = outputs[0].device
outputs = gather(outputs, target_device=target_dev)
else:
loss = criterion(outputs, target) # + kld.mean()
if add_criterion is not None:
loss2 = add_criterion(inputs, outputs) * weight
loss += loss2
inter, union = miou_class.get_iou(outputs, target)
inter_meter.update(inter)
union_meter.update(union)
loss = (loss - b).abs() + b
losses.update(loss.item(), inputs.size(0))
if add_criterion is not None:
nid_losses.update(loss2.item(), 1)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0: # print after every 100 batches
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
print_log_message(
"Epoch: %d[%d/%d]\t\tBatch Time:%.4f\t\tLoss:%.4f\t\tmiou:%.4f\t\tNID loss:%.4f"
% (epoch, i, len(dataset_loader), batch_time.avg, losses.avg,
miou, nid_losses.avg))
iou = inter_meter.sum / (union_meter.sum + 1e-10)
if greenhouse_use_trav:
miou = iou.mean() * 100
else:
miou = iou[[1, 2, 3]].mean() * 100
# miou = iou.mean() * 100
return iou, losses.avg
def validate(data_loader, model, criterion, device, epoch):
model.eval()
val_loss = AverageMeter()
val_cl_loss = AverageMeter()
val_loc_loss = AverageMeter()
batch_time = AverageMeter()
end = time.time()
with torch.no_grad():
for batch_idx, (images, boxes, labels) in enumerate(data_loader):
images = images.to(device)
boxes = boxes.to(device)
labels = labels.to(device)
confidences, locations = model(images)
regression_loss, classification_loss = criterion(
confidences, locations, labels, boxes)
loss = regression_loss + classification_loss
N = images.size(0)
val_loss.update(loss.item(), N)
val_cl_loss.update(classification_loss.item(), N)
val_loc_loss.update(regression_loss.item(), N)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx and batch_idx % 10 == 0:
print_log_message(
"Epoch: %d[%d/%d]\t\tBatch Time:%.4f\t\tCls Loss: %.4f(%.4f)\t \tLoc Loss: %.4f(%.4f)\t\tTotal Loss: %.4f(%.4f)"
% (epoch, batch_idx, len(data_loader), batch_time.avg,
val_cl_loss.val, val_cl_loss.avg, val_loc_loss.val,
val_loc_loss.avg, val_loss.val, val_loss.avg))
return val_loss.avg, val_cl_loss.avg, val_loc_loss.avg
def val_seg_ue(model,
dataset_loader,
criterion=None,
num_classes=21,
device='cuda',
use_depth=False,
add_criterion=None,
greenhouse_use_trav=False):
model.eval()
inter_meter = AverageMeter()
union_meter = AverageMeter()
batch_time = AverageMeter()
end = time.time()
miou_class = MIOU(num_classes=num_classes - 1)
if criterion:
losses = AverageMeter()
with torch.no_grad():
for i, batch in enumerate(dataset_loader):
inputs = batch[0].to(device=device)
target = batch[1].to(device=device)
if use_depth:
depth = batch[2].to(device=device)
outputs = model(inputs, depth)
else:
outputs = model(inputs)
if isinstance(outputs, OrderedDict):
out_aux = outputs['aux']
outputs = outputs['out']
else:
out_aux = outputs[1]
outputs = outputs[0]
outputs = outputs + 0.5 * out_aux
if criterion:
if device == 'cuda':
loss = criterion(outputs, target).mean()
if add_criterion is not None:
loss += add_criterion(inputs, outputs)
if isinstance(outputs, (list, tuple)):
target_dev = outputs[0].device
outputs = gather(outputs, target_device=target_dev)
else:
loss = criterion(outputs, target)
if add_criterion is not None:
loss += add_criterion(inputs, outputs)
losses.update(loss.item(), inputs.size(0))
inter, union = miou_class.get_iou(outputs, target)
inter_meter.update(inter)
union_meter.update(union)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0: # print after every 100 batches
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
loss_ = losses.avg if criterion is not None else 0
print_log_message(
"[%d/%d]\t\tBatch Time:%.4f\t\tLoss:%.4f\t\tmiou:%.4f" %
(i, len(dataset_loader), batch_time.avg, loss_, miou))
iou = inter_meter.sum / (union_meter.sum + 1e-10)
if greenhouse_use_trav:
miou = iou.mean() * 100
else:
# miou = np.array(iou)[1, 2, 3].mean() * 100
miou = iou[[1, 2, 3]].mean() * 100
# miou = iou.mean() * 100
print_info_message('Mean IoU: {0:.2f}'.format(miou))
if criterion:
return iou, losses.avg
else:
return iou, 0
def val_seg(model,
dataset_loader,
criterion=None,
num_classes=21,
device='cuda'):
model.eval()
inter_meter = AverageMeter()
union_meter = AverageMeter()
batch_time = AverageMeter()
end = time.time()
miou_class = MIOU(num_classes=num_classes)
if criterion:
losses = AverageMeter()
with torch.no_grad():
for i, (inputs, target) in enumerate(dataset_loader):
inputs = inputs.to(device=device)
target = target.to(device=device)
outputs = model(inputs)
if criterion:
if device == 'cuda':
loss = criterion(outputs, target).mean()
if isinstance(outputs, (list, tuple)):
target_dev = outputs[0].device
outputs = gather(outputs, target_device=target_dev)
else:
loss = criterion(outputs, target)
losses.update(loss.item(), inputs.size(0))
inter, union = miou_class.get_iou(outputs, target)
inter_meter.update(inter)
union_meter.update(union)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % 10 == 0: # print after every 100 batches
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
loss_ = losses.avg if criterion is not None else 0
print_log_message(
"[%d/%d]\t\tBatch Time:%.4f\t\tLoss:%.4f\t\tmiou:%.4f" %
(i, len(dataset_loader), batch_time.avg, loss_, miou))
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
print_info_message('Mean IoU: {0:.2f}'.format(miou))
if criterion:
return miou, losses.avg
else:
return miou, 0
def train(trainloader,
model,
criterion,
device,
optimizer,
class_encoding,
writer_idx,
class_weights=None,
writer=None):
"""Create the model and start the training."""
epoch_loss = 0
# For logging the training status
losses = AverageMeter()
nid_losses = AverageMeter()
kld_losses = AverageMeter()
batch_time = AverageMeter()
inter_meter = AverageMeter()
union_meter = AverageMeter()
miou_class = MIOU(num_classes=4)
model.train()
kld_layer = PixelwiseKLD()
with tqdm(total=len(trainloader)) as pbar:
for i_iter, batch in enumerate(tqdm(trainloader)):
images = batch[0].to(device)
labels = batch[1].to(device)
if args.use_depth:
depths = batch[2].to(device)
optimizer.zero_grad()
# Upsample the output of the model to the input size
# TODO: Change the input according to the model type
interp = None
if interp is not None:
if args.use_depth:
if args.model == 'espdnet':
pred = interp(model(images, depths))
elif args.model == 'espdnetue':
(pred, pred_aux) = interp(model(images, depths))
else:
if args.model == 'espdnet':
pred = interp(model(images))
elif args.model == 'espdnetue':
(pred, pred_aux) = interp(model(images))
elif args.model == 'deeplabv3':
output = model(images)
pred = interp(output['out'])
pred_aux = interp(output['aux'])
else:
if args.use_depth:
if args.model == 'espdnet':
pred = model(images, depths)
elif args.model == 'espdnetue':
(pred, pred_aux) = model(images, depths)
else:
if args.model == 'espdnet':
pred = model(images)
elif args.model == 'espdnetue':
(pred, pred_aux) = model(images)
elif args.model == 'deeplabv3':
output = model(images)
pred = output['out']
pred_aux = output['aux']
# print(pred.size())
# Model regularizer
kld = kld_layer(pred, pred_aux)
kld_losses.update(kld.mean().item(), 1)
if args.use_uncertainty:
loss = criterion(pred + 0.5 * pred_aux, labels,
kld) * 20 + kld.mean()
else:
loss = criterion(pred + 0.5 * pred_aux, labels) # + kld.mean()
inter, union = miou_class.get_iou(pred, labels)
inter_meter.update(inter)
union_meter.update(union)
losses.update(loss.item(), images.size(0))
# Optimise
loss.backward()
optimizer.step()
iou = inter_meter.sum / (union_meter.sum + 1e-10)
miou = iou.mean() * 100
# Write summary
writer.add_scalar('traversability_mask/train/loss', losses.avg, writer_idx)
writer.add_scalar('traversability_mask/train/nid_loss', nid_losses.avg,
writer_idx)
writer.add_scalar('traversability_mask/train/mean_IoU', miou, writer_idx)
writer.add_scalar('traversability_mask/train/traversable_plant_IoU',
iou[0], writer_idx)
writer.add_scalar('traversability_mask/train/other_plant_mean_IoU', iou[1],
writer_idx)
writer.add_scalar('traversability_mask/train/artificial_object_mean_IoU',
iou[2], writer_idx)
writer.add_scalar('traversability_mask/train/ground_mean_IoU', iou[3],
writer_idx)
writer.add_scalar('traversability_mask/train/learning_rate',
optimizer.param_groups[0]['lr'], writer_idx)
# writer.add_scalar('uest/train/kld', kld_losses.avg, writer_idx)
#
# Investigation of labels
#
# Before label conversion
if args.use_depth:
# model, images, depths=None, labels=None, predictions=None, class_encoding=None, writer=None, epoch=None, data=None, device=None
in_training_visualization_img(model,
images=images,
depths=depths,
labels=labels.long(),
class_encoding=class_encoding,
writer=writer,
epoch=writer_idx,
data='traversability_mask/train',
device=device)
else:
in_training_visualization_img(model,
images=images,
labels=labels.long(),
class_encoding=class_encoding,
writer=writer,
epoch=writer_idx,
data='traversability_mask/train',
device=device)
writer_idx += 1
print('taking snapshot ...')
return writer_idx
def main(args):
crop_size = args.crop_size
assert isinstance(crop_size, tuple)
print_info_message(
'Running Model at image resolution {}x{} with batch size {}'.format(
crop_size[0], crop_size[1], args.batch_size))
if not os.path.isdir(args.savedir):
os.makedirs(args.savedir)
num_gpus = torch.cuda.device_count()
device = 'cuda' if num_gpus > 0 else 'cpu'
if args.dataset == 'pascal':
from data_loader.segmentation.voc import VOCSegmentation, VOC_CLASS_LIST
train_dataset = VOCSegmentation(root=args.data_path,
train=True,
crop_size=crop_size,
scale=args.scale,
coco_root_dir=args.coco_path)
val_dataset = VOCSegmentation(root=args.data_path,
train=False,
crop_size=crop_size,
scale=args.scale)
seg_classes = len(VOC_CLASS_LIST)
class_wts = torch.ones(seg_classes)
elif args.dataset == 'city':
from data_loader.segmentation.cityscapes import CityscapesSegmentation, CITYSCAPE_CLASS_LIST
train_dataset = CityscapesSegmentation(root=args.data_path,
train=True,
size=crop_size,
scale=args.scale,
coarse=args.coarse)
val_dataset = CityscapesSegmentation(root=args.data_path,
train=False,
size=crop_size,
scale=args.scale,
coarse=False)
seg_classes = len(CITYSCAPE_CLASS_LIST)
class_wts = torch.ones(seg_classes)
class_wts[0] = 2.8149201869965
class_wts[1] = 6.9850029945374
class_wts[2] = 3.7890393733978
class_wts[3] = 9.9428062438965
class_wts[4] = 9.7702074050903
class_wts[5] = 9.5110931396484
class_wts[6] = 10.311357498169
class_wts[7] = 10.026463508606
class_wts[8] = 4.6323022842407
class_wts[9] = 9.5608062744141
class_wts[10] = 7.8698215484619
class_wts[11] = 9.5168733596802
class_wts[12] = 10.373730659485
class_wts[13] = 6.6616044044495
class_wts[14] = 10.260489463806
class_wts[15] = 10.287888526917
class_wts[16] = 10.289801597595
class_wts[17] = 10.405355453491
class_wts[18] = 10.138095855713
class_wts[19] = 0.0
elif args.dataset == 'greenhouse':
print(args.use_depth)
from data_loader.segmentation.greenhouse import GreenhouseRGBDSegmentation, GreenhouseDepth, GREENHOUSE_CLASS_LIST
train_dataset = GreenhouseDepth(root=args.data_path,
list_name='train_depth_ae.txt',
train=True,
size=crop_size,
scale=args.scale,
use_filter=True)
val_dataset = GreenhouseRGBDSegmentation(root=args.data_path,
list_name='val_depth_ae.txt',
train=False,
size=crop_size,
scale=args.scale,
use_depth=True)
class_weights = np.load('class_weights.npy')[:4]
print(class_weights)
class_wts = torch.from_numpy(class_weights).float().to(device)
seg_classes = len(GREENHOUSE_CLASS_LIST)
else:
print_error_message('Dataset: {} not yet supported'.format(
args.dataset))
exit(-1)
print_info_message('Training samples: {}'.format(len(train_dataset)))
print_info_message('Validation samples: {}'.format(len(val_dataset)))
from model.autoencoder.depth_autoencoder import espnetv2_autoenc
args.classes = 3
model = espnetv2_autoenc(args)
train_params = [{
'params': model.get_basenet_params(),
'lr': args.lr * args.lr_mult
}]
optimizer = optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay)
num_params = model_parameters(model)
flops = compute_flops(model,
input=torch.Tensor(1, 1, crop_size[0], crop_size[1]))
print_info_message(
'FLOPs for an input of size {}x{}: {:.2f} million'.format(
crop_size[0], crop_size[1], flops))
print_info_message('Network Parameters: {:.2f} million'.format(num_params))
writer = SummaryWriter(log_dir=args.savedir,
comment='Training and Validation logs')
try:
writer.add_graph(model,
input_to_model=torch.Tensor(1, 3, crop_size[0],
crop_size[1]))
except:
print_log_message(
"Not able to generate the graph. Likely because your model is not supported by ONNX"
)
start_epoch = 0
print('device : ' + device)
#criterion = nn.CrossEntropyLoss(weight=class_wts, reduction='none', ignore_index=args.ignore_idx)
#criterion = SegmentationLoss(n_classes=seg_classes, loss_type=args.loss_type,
# device=device, ignore_idx=args.ignore_idx,
# class_wts=class_wts.to(device))
criterion = nn.MSELoss()
# criterion = nn.L1Loss()
if num_gpus >= 1:
if num_gpus == 1:
# for a single GPU, we do not need DataParallel wrapper for Criteria.
# So, falling back to its internal wrapper
from torch.nn.parallel import DataParallel
model = DataParallel(model)
model = model.cuda()
criterion = criterion.cuda()
else:
from utilities.parallel_wrapper import DataParallelModel, DataParallelCriteria
model = DataParallelModel(model)
model = model.cuda()
criterion = DataParallelCriteria(criterion)
criterion = criterion.cuda()
if torch.backends.cudnn.is_available():
import torch.backends.cudnn as cudnn
cudnn.benchmark = True
cudnn.deterministic = True
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=args.workers)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=args.workers)
if args.scheduler == 'fixed':
step_size = args.step_size
step_sizes = [
step_size * i
for i in range(1, int(math.ceil(args.epochs / step_size)))
]
from utilities.lr_scheduler import FixedMultiStepLR
lr_scheduler = FixedMultiStepLR(base_lr=args.lr,
steps=step_sizes,
gamma=args.lr_decay)
elif args.scheduler == 'clr':
step_size = args.step_size
step_sizes = [
step_size * i
for i in range(1, int(math.ceil(args.epochs / step_size)))
]
from utilities.lr_scheduler import CyclicLR
lr_scheduler = CyclicLR(min_lr=args.lr,
cycle_len=5,
steps=step_sizes,
gamma=args.lr_decay)
elif args.scheduler == 'poly':
from utilities.lr_scheduler import PolyLR
lr_scheduler = PolyLR(base_lr=args.lr,
max_epochs=args.epochs,
power=args.power)
elif args.scheduler == 'hybrid':
from utilities.lr_scheduler import HybirdLR
lr_scheduler = HybirdLR(base_lr=args.lr,
max_epochs=args.epochs,
clr_max=args.clr_max,
cycle_len=args.cycle_len)
elif args.scheduler == 'linear':
from utilities.lr_scheduler import LinearLR
lr_scheduler = LinearLR(base_lr=args.lr, max_epochs=args.epochs)
else:
print_error_message('{} scheduler Not supported'.format(
args.scheduler))
exit()
print_info_message(lr_scheduler)
with open(args.savedir + os.sep + 'arguments.json', 'w') as outfile:
import json
arg_dict = vars(args)
arg_dict['model_params'] = '{} '.format(num_params)
arg_dict['flops'] = '{} '.format(flops)
json.dump(arg_dict, outfile)
extra_info_ckpt = '{}_{}_{}'.format(args.model, args.s, crop_size[0])
best_loss = 0.0
for epoch in range(start_epoch, args.epochs):
lr_base = lr_scheduler.step(epoch)
# set the optimizer with the learning rate
# This can be done inside the MyLRScheduler
lr_seg = lr_base * args.lr_mult
optimizer.param_groups[0]['lr'] = lr_seg
# optimizer.param_groups[1]['lr'] = lr_seg
# Train
model.train()
losses = AverageMeter()
for i, batch in enumerate(train_loader):
inputs = batch[1].to(device=device) # Depth
target = batch[0].to(device=device) # RGB
outputs = model(inputs)
if device == 'cuda':
loss = criterion(outputs, target).mean()
if isinstance(outputs, (list, tuple)):
target_dev = outputs[0].device
outputs = gather(outputs, target_device=target_dev)
else:
loss = criterion(outputs, target)
losses.update(loss.item(), inputs.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
# if not (i % 10):
# print("Step {}, write images".format(i))
# image_grid = torchvision.utils.make_grid(outputs.data.cpu()).numpy()
# writer.add_image('Autoencoder/results/train', image_grid, len(train_loader) * epoch + i)
writer.add_scalar('Autoencoder/Loss/train', loss.item(),
len(train_loader) * epoch + i)
print_info_message('Running batch {}/{} of epoch {}'.format(
i + 1, len(train_loader), epoch + 1))
train_loss = losses.avg
writer.add_scalar('Autoencoder/LR/seg', round(lr_seg, 6), epoch)
# Val
if epoch % 5 == 0:
losses = AverageMeter()
with torch.no_grad():
for i, batch in enumerate(val_loader):
inputs = batch[2].to(device=device) # Depth
target = batch[0].to(device=device) # RGB
outputs = model(inputs)
if device == 'cuda':
loss = criterion(outputs, target) # .mean()
if isinstance(outputs, (list, tuple)):
target_dev = outputs[0].device
outputs = gather(outputs, target_device=target_dev)
else:
loss = criterion(outputs, target)
losses.update(loss.item(), inputs.size(0))
image_grid = torchvision.utils.make_grid(
outputs.data.cpu()).numpy()
writer.add_image('Autoencoder/results/val', image_grid,
epoch)
image_grid = torchvision.utils.make_grid(
inputs.data.cpu()).numpy()
writer.add_image('Autoencoder/inputs/val', image_grid,
epoch)
image_grid = torchvision.utils.make_grid(
target.data.cpu()).numpy()
writer.add_image('Autoencoder/target/val', image_grid,
epoch)
val_loss = losses.avg
print_info_message(
'Running epoch {} with learning rates: base_net {:.6f}, segment_net {:.6f}'
.format(epoch, lr_base, lr_seg))
# remember best miou and save checkpoint
is_best = val_loss < best_loss
best_loss = min(val_loss, best_loss)
weights_dict = model.module.state_dict(
) if device == 'cuda' else model.state_dict()
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.model,
'state_dict': weights_dict,
'best_loss': best_loss,
'optimizer': optimizer.state_dict(),
}, is_best, args.savedir, extra_info_ckpt)
writer.add_scalar('Autoencoder/Loss/val', val_loss, epoch)
writer.close()
def test(testloader,
prob_model,
likelihood,
seg_model,
V,
device,
writer,
feature_mean=0.0,
mask_mean=0.0,
mode='test',
max_image=100):
"""Create the model and start the evaluation process."""
# For logging the training status
sigmoid = nn.Sigmoid()
prob_sum_meter = AverageMeter()
losses = AverageMeter()
## model for evaluation
seg_model.to(device)
seg_model.eval()
prob_model.to(device)
prob_model.eval()
# TODO: Change this (implement the same function in 'utility/utils.py', or uncomment the code below with slight modification)
with torch.no_grad(), gpytorch.settings.fast_pred_var():
# Calculate a constant c
for i_iter, batch in enumerate(tqdm(testloader)):
if i_iter == max_image:
break
images = batch["rgb"].to(device)
images_orig = batch["rgb_orig"].to(device)
masks = batch["mask"].to(device)
if args.use_depth:
depths = batch["depth"].to(device)
output_dict = get_output(seg_model, images)
feature = output_dict['feature']
f_orig_size = feature.size()
# Downsample features
feature = F.interpolate(feature, (128, 240), mode='nearest')
f_ds_size = feature.size()
# Downsample masks
ih = torch.linspace(0, masks.size(1) - 1, 128).long()
iw = torch.linspace(0, masks.size(2) - 1, 240).long()
masks = masks[:, ih[:, None], iw]
# masks = F.interpolate(masks, (8,15), mode='nearest')
# masks = torch.squeeze(masks)
feature = feature.transpose(1, 2).transpose(2, 3).reshape(
(-1, feature.size(1)))
# Get prediction result
observed_pred = likelihood(prob_model(feature))
prob = observed_pred.mean
prob_int = prob.clone()
prob_int[prob_int < 0] = 0.0
prob_int = prob_int.reshape((-1, 1, f_ds_size[2], f_ds_size[3]))
prob_int = F.interpolate(prob_int,
(f_orig_size[2], f_orig_size[3]),
mode='bilinear')
if i_iter == 0:
image_tensor = images_orig
mask_tensor = masks
prob_tensor = prob_int
else:
image_tensor = torch.cat((image_tensor, images_orig))
mask_tensor = torch.cat((mask_tensor, masks))
prob_tensor = torch.cat((prob_tensor, prob_int))
mask_tensor = torch.reshape(mask_tensor,
(mask_tensor.size(0), -1,
mask_tensor.size(1), mask_tensor.size(2)))
image_grid = torchvision.utils.make_grid(
image_tensor.data.cpu()).numpy()
mask_grid = torchvision.utils.make_grid(mask_tensor.data.cpu()).numpy()
prob_grid = torchvision.utils.make_grid(prob_tensor.data.cpu()).numpy()
writer.add_image('traversability_mask/{}/image'.format(mode),
image_grid, i_iter)
writer.add_image('traversability_mask/{}/mask'.format(mode), mask_grid,
i_iter)
writer.add_image('traversability_mask/{}/prob'.format(mode), prob_grid,
i_iter)
feature += feature_mean
prob += mask_mean
visualize_gp_function(feature, prob,
masks.squeeze().flatten(), V, writer, mode)
