class Trainer(object):
def __init__(self, cfg, output_dir):
self.cfg = cfg
self.output_dir = output_dir
self.logger = logging.getLogger("MaskPyramid")
self.tbSummary = TensorboardSummary(output_dir)
self.writer = self.tbSummary.create_summary()
# self.model = MaskPyramids(cfg)
self.model = HighResolutionNet(cfg)
self.model.init_weights(cfg.MODEL.PRETRAINED)
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(cfg)
train_params = [{'params': self.model.parameters(), 'lr': cfg.SOLVER.BASE_LR}]
self.optimizer = torch.optim.SGD(train_params, momentum=cfg.SOLVER.MOMENTUM,
weight_decay=cfg.SOLVER.WEIGHT_DECAY, nesterov=cfg.SOLVER.NESTEROV)
self.evaluator = Evaluator(self.nclass)
self.scheduler = LR_Scheduler(cfg.SOLVER.SCHEDULE_TYPE, cfg.SOLVER.BASE_LR,
cfg.SOLVER.EPOCHES, len(self.train_loader))
self.start_epoch = 0
self.best_pred = 0.0
self.meters = {'start_time': time.time(),
'total_iters': cfg.SOLVER.EPOCHES*len(self.train_loader)}
# log_gpu_stat(self.logger)
def load_weights(self, path=None, subdict='model', continue_train=False):
state_dict = torch.load(path if path else self.cfg.MODEL.WEIGHT)
if subdict == 'hrnet':
# model_state_file = config.TEST.MODEL_FILE
pretrained_dict = torch.load(path)
if list(pretrained_dict.keys())[0].startswith('model.'):
model_dict = self.model.state_dict()
pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items()
if k[6:] in model_dict.keys()}
# for k, _ in pretrained_dict.items():
# self.logger.info(
# '=> loading {} from pretrained model'.format(k))
model_dict.update(pretrained_dict)
self.model.load_state_dict(model_dict)
# model_dict = self.model.state_dict()
elif subdict == 'model':
weights = state_dict[subdict]
self.model.load_state_dict(weights)
if self.optimizer is not None and 'optimizer' in state_dict.keys():
self.optimizer.load_state_dict(state_dict["optimizer"])
if 'best_pred' in state_dict.keys():
self.best_pred = state_dict["best_pred"]
if continue_train and 'epoch' in state_dict.keys():
self.start_epoch = state_dict["epoch"]
else:
self.model.load_state_dict(state_dict)
def mask_inst_img(self, images_np, sematic_a_image_np, output, idx, level=-1, label_mode=1, insts_a_image=None):
pyramids = output['all_pyramids'][idx]
inst_target = None
if level>=0:
import pdb; pdb.set_trace()
bg_and_masks = [F.interpolate(pyramids[0].get_mask(3)[:,[0]], 513, mode='bilinear', align_corners=True)]
pred_labels = [0]
sematic_out = output['sematic_out'][idx]
for pyr in pyramids:
# import pdb; pdb.set_trace()
pos = tuple(pyr.pos.tolist())
largest_mask = pyr.get_mask(3)
mask_513 = F.interpolate(largest_mask[:,[1]], 513, mode='bilinear', align_corners=True)
# mask_513_map = mask_513[0].max(0)[1]
bg_and_masks.append(mask_513)
# prepare label
scaled_sematic_out = F.interpolate(sematic_out[None], pyr.init_size, mode='bilinear', align_corners=True)
label = scaled_sematic_out[0].max(0)[1][pos[0], pos[1]].item()
pred_labels.append(label)
# import pdb; pdb.set_trace()
bg_and_masks_np = torch.cat(bg_and_masks).squeeze(1).max(0)[1].detach().cpu().numpy()
inst_output = self.rend_on_image(images_np, bg_and_masks_np, pred_labels)
return inst_output, inst_target
def rend_on_image_v8(self, image_np, masks_np, labels_list):
# label 从0开始，0 表示unlabeled
alpha = 0.5
color_bias = 1
# import pdb; pdb.set_trace()
# print('masks_np unique',np.unique(masks_np), 'labels_list', labels_list)
class_count = [0 for _ in range(20)]
colors = np.array([
# [0, 100, 0],
[128, 64, 128],
[244, 35, 232],
[70, 70, 70],
[102, 102, 156],
[190, 153, 153],
[153, 153, 153],
[250, 170, 30],
[220, 220, 0],
[107, 142, 35],
[152, 251, 152],
[0, 130, 180],
[220, 20, 60],
[255, 0, 0],
[0, 0, 142],
[0, 0, 70],
[0, 60, 100],
[0, 80, 100],
[0, 0, 230],
[119, 11, 32]])
# import pdb; pdb.set_trace()
masked = image_np.copy()
# import pdb; pdb.set_trace()
for i, label in enumerate(labels_list):
if label == -1:
continue
mask_idx = np.nonzero(masks_np == i)
masked[mask_idx[0], mask_idx[1], :] *= 1.0 - alpha
color = (colors[label]+class_count[label]*color_bias)/255.0
# try:
# color = (colors[label]+class_count[label]*color_bias)/255.0
# except:
# import pdb; pdb.set_trace()
masked[mask_idx[0], mask_idx[1], :] += alpha * color
# import pdb; pdb.set_trace()
contours, hierarchy = cv2.findContours(
(masks_np == i).astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
)
# masked = cv2.drawContours(masked, contours, -1, (1.,1.,1.), -1)
masked = cv2.drawContours(masked, contours, -1, (1.,1.,1.), 1)
class_count[label] += 1
return masked
def show_image_v8(self, images, output, targets):
N,C,H,W = images.shape
masked_imgs = []
origin_image_ON = True
ori_sematic_target_ON = False
ori_instance_target_ON = True
car_only_ON = False
pred_sematic_ON = True
pred_instance_ON = True
pred_stat_ON = True
for i, image in enumerate(images):
parts_to_show = []
images_np = image.permute(1,2,0).detach().cpu().numpy()
images_np = ((images_np - images_np.min()) / (images_np.max() - images_np.min()))[...,::-1]
if origin_image_ON:
parts_to_show.append(images_np)
semat_a_target, insts_a_target = targets['label'][i], targets['instance'][i]
# targets prepared in forwards
# out_one = output['in_dict_watches'][i]
# target_fwd_label = (out_one['target_levels']['labels']+1).tolist()
sematic_target = output['targets']['label'][i]
instacne_target = output['targets']['instance'][i]
# sematic_out_513 = output['sematic_out'][6]
sematic_out_513 = output['sematic_out']
if ori_sematic_target_ON:
sematic_target_np = sematic_target.detach().cpu().numpy()
ori_sematic_target = self.rend_on_image_v8(images_np, sematic_target_np, range(19))
parts_to_show.append(ori_sematic_target)
if ori_instance_target_ON:
instacne_target_np = instacne_target.detach().cpu().numpy()
ins_cats = [sematic_target[instacne_target == ins_idx].unique().item() for ins_idx in instacne_target.unique()]
if car_only_ON:
ins_cats = [13 if label==13 else -1 for label in ins_cats]
ori_instacne_target = self.rend_on_image_v8(images_np, instacne_target_np, ins_cats)
parts_to_show.append(ori_instacne_target)
# import pdb; pdb.set_trace()
# prediction - sematic
if pred_sematic_ON:
sematic_predict_np = sematic_out_513[i].max(0)[1].detach().cpu().numpy()
sematic_predict = self.rend_on_image_v8(images_np, sematic_predict_np, range(19))
parts_to_show.append(sematic_predict)
if not self.cfg.SOLVER.SEMATIC_ONLY:
parymids = output['ins_pyramids']
# prediction - instance
if pred_instance_ON:
# import pdb; pdb.set_trace()
instance_predict = torch.cat([pyr.get_mask(513) for pyr in parymids[i]], dim=1) if len(parymids[i]) else torch.empty((1,1,513,513))
instance_predict_513 = F.interpolate(instance_predict, 513, mode='bilinear', align_corners=True)
instance_predict_np = instance_predict_513[0].max(0)[1].detach().cpu().numpy()
ins_cats_predict = [pyr.tar_cat for pyr in parymids[i]]
# import pdb; pdb.set_trace()
pred_instance = self.rend_on_image_v8(images_np, instance_predict_np, ins_cats_predict)
# pred_instance = self.rend_on_image_v8(images_np, instance_predict_np, range(19))
if pred_stat_ON:
cv2.putText(pred_instance,'#:{}'.format(len(parymids[i])),(30,30),cv2.FONT_HERSHEY_COMPLEX,1,(0,0,255),1)
for pyr in parymids[i]:
# import pdb; pdb.set_trace()
pyr_pos_513 = pyr.get_rel_pos(7)[::-1]
cv2.circle(pred_instance, pyr_pos_513, 3, (0,0,255), -1)
parts_to_show.append(pred_instance)
masked_imgs.append(np.hstack(parts_to_show))
masked_show = np.vstack(masked_imgs)
cv2.imshow('Observation V5', masked_show)
cv2.waitKey(10)
def training(self, epoch):
train_loss = 0.0
self.model.train()
end = time.time()
# tbar = tqdm(self.train_loader)
# for i, sample in enumerate(tbar):
for i, sample in enumerate(self.train_loader):
# image = sample['image'].to(self.device)
# target = {'label': sample['label'].to(self.device),
# 'instance': sample['instance'].to(self.device)}
image, label, instance = sample
image, label, instance = image.to(self.device), label.to(self.device), instance.to(self.device)
target = {'label': label, 'instance': instance}
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output_dict = self.model(image, target)
# try:
# output_dict = self.model(image, target)
# except:
# log_gpu_stat(self.logger)
# print('Num of instances:', self.model.log_dict['InstPyr_inst_count'])
# exit()
# import pdb; pdb.set_trace()
loss_dict = output_dict['loss_dict']
losses = sum(loss for loss in loss_dict.values())
# import pdb; pdb.set_trace()
losses.backward()
self.optimizer.step()
batch_time = time.time() - end
end = time.time()
# self.model.log_dict['loss_dict'] = loss_dict
train_loss += losses.item()
if i % 20 == 0 or i == len(self.train_loader) -1:
curr_iter = epoch*len(self.train_loader)+ i + 1
sepent_time = time.time() - self.meters['start_time']
eta_seconds = sepent_time * self.meters['total_iters'] / curr_iter - sepent_time
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
if self.cfg.SOLVER.SEMATIC_ONLY:
self.logger.info(('Ep:{}/{}|Iter:{}/{}|Eta:{}|SematicLoss:{:2.4}|Class:{:2.4}').format(
epoch, self.cfg.SOLVER.EPOCHES, i, len(self.train_loader),
eta_string, losses.item(), loss_dict['sematic'].item()))
else:
self.logger.info(('Ep:{}/{}|Iter:{}/{}|Eta:{}|SematicLoss:{:2.4}|Class:{:2.4}').format(
epoch, self.cfg.SOLVER.EPOCHES, i, len(self.train_loader),
eta_string, losses.item(), loss_dict['sematic'].item()))
# self.logger.info(('Ep:{}/{}|Iter:{}/{}|Eta:{}|Loss:{:2.4}|Class:{:2.4}|'+\
# 'L0:{:2.4}|L1:{:2.4}|L2:{:2.4}|L3:{:2.4}|PyrNum:{}|#pyr0:{}|'+\
# '#pyr1:{}|#pyr2:{}|#pyr3:{}|').format(
# epoch, self.cfg.SOLVER.EPOCHES, i, len(self.train_loader),
# eta_string, losses.item(), loss_dict['class_loss'].item(), loss_dict['level_0'].item(),
# loss_dict['level_1'].item(), loss_dict['level_2'].item(), loss_dict['level_3'].item(),
# self.model.log_dict['InstPyr_inst_count'], self.model.log_dict['pyr_num_l0'],
# self.model.log_dict['pyr_num_l1'], self.model.log_dict['pyr_num_l2'],
# self.model.log_dict['pyr_num_l3'],
# ))
if self.cfg.SOLVER.SHOW_IMAGE and i % 10 == 0:
self.show_image_v8(image, output_dict, target)
self.writer.add_scalar('train/loss_epoch', train_loss, epoch)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
test_loss = 0.0
tbar = tqdm(self.val_loader, desc='\r')
for i, sample in enumerate(tbar):
image, label, instance = sample
image, label, instance = image.to(self.device), label.to(self.device), instance.to(self.device)
target = {'label': label, 'instance': instance}
with torch.no_grad():
output = self.model(image, target)
# output = self.model(image)
# import pdb; pdb.set_trace()
sematic_out = output['sematic_out']
pred = sematic_out.data.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(label.detach().cpu().numpy(), pred)
if self.cfg.SOLVER.SHOW_IMAGE and i % 10 == 0:
self.show_image_v8(image, output, target)
# Fast test during the training
# import pdb; pdb.set_trace()
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
self.logger.info('Evalueat report: mIoU: {:3.4}| Acc: {:3.4}| Acc_class: {:3.4}| fwIoU: {:3.4}| previousBest: {:3.4}|'.format(
mIoU, Acc, Acc_class, FWIoU, float(self.best_pred)
))
if mIoU > self.best_pred:
is_best = True
self.best_pred = mIoU
save_data = {}
save_data["epoch"] = epoch + 1
save_data["best_pred"] = self.best_pred
save_data["model"] = self.model.state_dict()
if self.optimizer is not None:
save_data["optimizer"] = self.optimizer.state_dict()
if self.scheduler is not None:
save_data["scheduler"] = self.scheduler.__dict__
torch.save(save_data, self.output_dir+"/model_Epoch_{}.pth".format(epoch))
class Trainer(object):
def __init__(self, cfg, output_dir):
self.cfg = cfg
self.output_dir = output_dir
self.logger = logging.getLogger("MaskPyramid")
self.tbSummary = TensorboardSummary(output_dir)
self.writer = self.tbSummary.create_summary()
self.model = MaskPyramids(cfg)
# self.optimizer = make_optimizer(cfg, self.model)
# self.scheduler = make_lr_scheduler(cfg, self.optimizer)
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
cfg)
train_params = [{
'params': self.model.parameters(),
'lr': cfg.SOLVER.BASE_LR
}]
self.optimizer = torch.optim.SGD(train_params,
momentum=cfg.SOLVER.MOMENTUM,
weight_decay=cfg.SOLVER.WEIGHT_DECAY,
nesterov=cfg.SOLVER.NESTEROV)
self.evaluator = Evaluator(self.nclass)
self.scheduler = LR_Scheduler(cfg.SOLVER.SCHEDULE_TYPE,
cfg.SOLVER.BASE_LR, cfg.SOLVER.EPOCHES,
len(self.train_loader))
self.start_epoch = 0
self.best_pred = 0.0
self.meters = {
'start_time': time.time(),
'total_iters': cfg.SOLVER.EPOCHES * len(self.train_loader)
}
def load_weights(self, path=None, subdict='model', continue_train=False):
state_dict = torch.load(path if path else self.cfg.MODEL.WEIGHT)
if subdict:
weights = state_dict[subdict]
self.model.load_state_dict(weights)
if self.optimizer is not None and 'optimizer' in state_dict.keys():
self.optimizer.load_state_dict(state_dict["optimizer"])
if 'best_pred' in state_dict.keys():
self.best_pred = state_dict["best_pred"]
if continue_train and 'epoch' in state_dict.keys():
self.start_epoch = state_dict["epoch"]
else:
self.model.load_state_dict(state_dict)
def training(self, epoch):
train_loss = 0.0
self.model.train()
end = time.time()
# tbar = tqdm(self.train_loader)
# for i, sample in enumerate(tbar):
for i, sample in enumerate(self.train_loader):
image = sample['image'].to(self.device)
target = {
'label': sample['label'].to(self.device),
'instance': sample['instance'].to(self.device)
}
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
loss_dict = self.model(image, target)
losses = sum(loss for loss in loss_dict.values())
self.model.log_dict['loss_dict'] = loss_dict
losses.backward()
self.optimizer.step()
batch_time = time.time() - end
end = time.time()
self.model.log_dict['loss_dict'] = loss_dict
train_loss += losses.item()
if i % 20 == 0 or i == len(self.train_loader) - 1:
curr_iter = epoch * len(self.train_loader) + i + 1
sepent_time = time.time() - self.meters['start_time']
eta_seconds = sepent_time * self.meters[
'total_iters'] / curr_iter - sepent_time
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
self.logger.info(('Ep:{}/{}|Iter:{}/{}|Eta:{}|Loss:{:2.4}|Class:{:2.4}|'+\
'L0:{:2.4}|L1:{:2.4}|L2:{:2.4}|L3:{:2.4}|PyrNum:{}|#pyr0:{}|'+\
'#pyr1:{}|#pyr2:{}|#pyr3:{}|').format(
epoch, self.cfg.SOLVER.EPOCHES, i, len(self.train_loader),
eta_string, losses.item(), loss_dict['class_loss'].item(), loss_dict['level_0'].item(),
loss_dict['level_1'].item(), loss_dict['level_2'].item(), loss_dict['level_3'].item(),
self.model.log_dict['InstPyr_inst_count'], self.model.log_dict['pyr_num_l0'],
self.model.log_dict['pyr_num_l1'], self.model.log_dict['pyr_num_l2'],
self.model.log_dict['pyr_num_l3'],
))
self.writer.add_scalar('train/loss_epoch', train_loss, epoch)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
test_loss = 0.0
tbar = tqdm(self.val_loader, desc='\r')
for i, sample in enumerate(tbar):
image = sample['image'].to(self.device)
target = {
'label': sample['label'].to(self.device),
'instance': sample['instance'].to(self.device)
}
with torch.no_grad():
output = self.model(image, target)
# import pdb; pdb.set_trace()
sematic_out = output['sematic_out']
pred = sematic_out.data.cpu().numpy()
target = sample['label'].numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
self.logger.info(
'Evalueat report: mIoU: {:3.4}| Acc: {:3.4}| Acc_class: {:3.4}| fwIoU: {:3.4}|'
.format(mIoU, Acc, Acc_class, FWIoU))
if mIoU > self.best_pred:
is_best = True
self.best_pred = mIoU
save_data = {}
save_data["epoch"] = epoch + 1
save_data["best_pred"] = self.best_pred
save_data["model"] = self.model.state_dict()
if self.optimizer is not None:
save_data["optimizer"] = self.optimizer.state_dict()
if self.scheduler is not None:
save_data["scheduler"] = self.scheduler.__dict__
torch.save(
save_data,
self.output_dir + "/model_Epoch_{:3d}.pth".format(epoch))
class Trainer(object):
def __init__(self, cfg, output_dir):
self.cfg = cfg
self.output_dir = output_dir
self.logger = logging.getLogger("MaskPyramid")
self.tbSummary = TensorboardSummary(output_dir)
self.writer = self.tbSummary.create_summary()
self.model = MaskPyramids(cfg)
# self.optimizer = make_optimizer(cfg, self.model)
# self.scheduler = make_lr_scheduler(cfg, self.optimizer)
self.device = torch.device(cfg.MODEL.DEVICE)
self.model.to(self.device)
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
cfg)
train_params = [{
'params': self.model.parameters(),
'lr': cfg.SOLVER.BASE_LR
}]
self.optimizer = torch.optim.SGD(train_params,
momentum=cfg.SOLVER.MOMENTUM,
weight_decay=cfg.SOLVER.WEIGHT_DECAY,
nesterov=cfg.SOLVER.NESTEROV)
self.evaluator = Evaluator(self.nclass)
self.scheduler = LR_Scheduler(cfg.SOLVER.SCHEDULE_TYPE,
cfg.SOLVER.BASE_LR, cfg.SOLVER.EPOCHES,
len(self.train_loader))
self.start_epoch = 0
self.best_pred = 0.0
self.meters = {
'start_time': time.time(),
'total_iters': cfg.SOLVER.EPOCHES * len(self.train_loader)
}
log_gpu_stat(self.logger)
def load_weights(self, path=None, subdict='model', continue_train=False):
state_dict = torch.load(path if path else self.cfg.MODEL.WEIGHT)
if subdict:
weights = state_dict[subdict]
self.model.load_state_dict(weights)
if self.optimizer is not None and 'optimizer' in state_dict.keys():
self.optimizer.load_state_dict(state_dict["optimizer"])
if 'best_pred' in state_dict.keys():
self.best_pred = state_dict["best_pred"]
if continue_train and 'epoch' in state_dict.keys():
self.start_epoch = state_dict["epoch"]
else:
self.model.load_state_dict(state_dict)
def rend_on_image(self, image_np, masks_np, labels):
alpha = 0.5
colors = np.array([[0, 100, 0], [128, 64, 128], [244, 35, 232],
[70, 70, 70], [102, 102, 156], [190, 153, 153],
[153, 153, 153], [250, 170, 30], [220, 220, 0],
[107, 142, 35], [152, 251, 152], [0, 130, 180],
[220, 20, 60], [255, 0, 0], [0, 0, 142], [0, 0, 70],
[0, 60, 100], [0, 80, 100], [0, 0, 230],
[119, 11, 32]])
# import pdb; pdb.set_trace()
masked = image_np.copy()
for i, label in enumerate(labels):
mask_idx = np.nonzero(masks_np == i)
masked[mask_idx[0], mask_idx[1], :] *= 1.0 - alpha
color = (colors[label + 1] + i) / 255.0
masked[mask_idx[0], mask_idx[1], :] += alpha * color
# import pdb; pdb.set_trace()
contours, hierarchy = cv2.findContours(
(masks_np == i).astype(np.uint8), cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)
# masked = cv2.drawContours(masked, contours, -1, (1.,1.,1.), -1)
masked = cv2.drawContours(masked, contours, -1, (1., 1., 1.), 1)
return masked, color
def show_image(self, images, output, targets):
N, C, H, W = images.shape
masked_imgs = []
for i, image in enumerate(images):
images_np = image.permute(1, 2, 0).detach().cpu().numpy()
images_np = ((images_np - images_np.min()) /
(images_np.max() - images_np.min()))[..., ::-1]
semt_a_image, insts_a_image = targets['label'][i], targets[
'instance'][i]
class_of_inst = [
semt_a_image[insts_a_image == j].unique().item()
for j in range(len(insts_a_image.unique()))
]
class_names = self.train_loader.dataset.class_names
insts_a_image_np = insts_a_image.detach().cpu().numpy()
masked_target, color1 = self.rend_on_image(images_np,
insts_a_image_np,
class_of_inst)
# prediction
# import pdb; pdb.set_trace()
sematic_a_image = output['sematic_out'][i]
sematic_a_image_np = sematic_a_image.max(
0)[1].detach().cpu().numpy()
masked_sematic, color2 = self.rend_on_image(
images_np, sematic_a_image_np, range(19))
masked_imgs.append(
np.hstack((images_np, masked_target, masked_sematic)))
masked_show = np.vstack(masked_imgs)
cv2.imshow('Observation', masked_show)
cv2.waitKey(1)
def training(self, epoch):
train_loss = 0.0
self.model.train()
end = time.time()
# tbar = tqdm(self.train_loader)
# for i, sample in enumerate(tbar):
for i, sample in enumerate(self.train_loader):
# image = sample['image'].to(self.device)
# target = {'label': sample['label'].to(self.device),
# 'instance': sample['instance'].to(self.device)}
image, label, instance = sample
image, label, instance = image.to(self.device), label.to(
self.device), instance.to(self.device)
target = {'label': label, 'instance': instance}
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output_dict = self.model(image, target)
loss_dict = output_dict['loss_dict']
losses = sum(loss for loss in loss_dict.values())
self.model.log_dict['loss_dict'] = loss_dict
losses.backward()
self.optimizer.step()
batch_time = time.time() - end
end = time.time()
self.model.log_dict['loss_dict'] = loss_dict
train_loss += losses.item()
if i % 20 == 0 or i == len(self.train_loader) - 1:
curr_iter = epoch * len(self.train_loader) + i + 1
sepent_time = time.time() - self.meters['start_time']
eta_seconds = sepent_time * self.meters[
'total_iters'] / curr_iter - sepent_time
eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
self.logger.info(('Ep:{}/{}|Iter:{}/{}|Eta:{}|Loss:{:2.4}|Class:{:2.4}|'+\
'L0:{:2.4}|L1:{:2.4}|L2:{:2.4}|L3:{:2.4}|PyrNum:{}|#pyr0:{}|'+\
'#pyr1:{}|#pyr2:{}|#pyr3:{}|').format(
epoch, self.cfg.SOLVER.EPOCHES, i, len(self.train_loader),
eta_string, losses.item(), loss_dict['class_loss'].item(), loss_dict['level_0'].item(),
loss_dict['level_1'].item(), loss_dict['level_2'].item(), loss_dict['level_3'].item(),
self.model.log_dict['InstPyr_inst_count'], self.model.log_dict['pyr_num_l0'],
self.model.log_dict['pyr_num_l1'], self.model.log_dict['pyr_num_l2'],
self.model.log_dict['pyr_num_l3'],
))
if self.cfg.SOLVER.SHOW_IMAGE and i % 2 == 0:
self.show_image(image, output_dict, target)
self.writer.add_scalar('train/loss_epoch', train_loss, epoch)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
test_loss = 0.0
tbar = tqdm(self.val_loader, desc='\r')
for i, sample in enumerate(tbar):
image, label, instance = sample
image, label, instance = image.to(self.device), label.to(
self.device), instance.to(self.device)
target = {'label': label, 'instance': instance}
with torch.no_grad():
output = self.model(image, target)
# import pdb; pdb.set_trace()
sematic_out = output['sematic_out']
pred = sematic_out.data.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(label.detach().cpu().numpy(), pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
self.logger.info(
'Evalueat report: mIoU: {:3.4}| Acc: {:3.4}| Acc_class: {:3.4}| fwIoU: {:3.4}|'
.format(mIoU, Acc, Acc_class, FWIoU))
if mIoU > self.best_pred:
is_best = True
self.best_pred = mIoU
save_data = {}
save_data["epoch"] = epoch + 1
save_data["best_pred"] = self.best_pred
save_data["model"] = self.model.state_dict()
if self.optimizer is not None:
save_data["optimizer"] = self.optimizer.state_dict()
if self.scheduler is not None:
save_data["scheduler"] = self.scheduler.__dict__
torch.save(
save_data,
self.output_dir + "/model_Epoch_{:3d}.pth".format(epoch))
