def validate(self): print 'Validating' ggnn_grid_size = self.opts['ggnn_grid_size'] self.model.ggnn.encoder.eval() self.model.temperature = 0 self.model.mode = "test" # Leave LSTM in train mode with torch.no_grad(): ious = [] for step, data in enumerate(tqdm(self.val_loader)): output = self.model(data['img'].to(device), data['fwd_poly'].to(device)) pred_polys = output['pred_polys'].data.numpy() # Get IoU iou = 0 orig_poly = data['orig_poly'] for i in range(pred_polys.shape[0]): p = pred_polys[i] mask_poly = utils.get_masked_poly( p, self.model.ggnn.ggnn_grid_size) mask_poly = utils.class_to_xy( mask_poly, self.model.ggnn.ggnn_grid_size) curr_gt_poly_112 = utils.poly01_to_poly0g( orig_poly[i], ggnn_grid_size) i, masks = metrics.iou_from_poly( np.array(mask_poly, dtype=np.int32), np.array(curr_gt_poly_112, dtype=np.int32), ggnn_grid_size, ggnn_grid_size) iou += i iou = iou / pred_polys.shape[0] ious.append(iou) del (output) del (pred_polys) iou = np.mean(ious) self.val_writer.add_scalar('iou', float(iou), self.global_step) print '[VAL] IoU: %f' % iou self.model.temperature = self.opts['temperature'] self.model.mode = "train_ggnn" self.model.ggnn.encoder.train()
def validate(self): print 'Validating' self.model.encoder.eval() self.model.first_v.eval() # Leave LSTM in train mode ious = [] accuracies = [] with torch.no_grad(): for step, data in enumerate(tqdm(self.val_loader)): output = self.model(data['img'].to(device), data['fwd_poly'].to(device)) # Get accuracy accuracy = metrics.train_accuracy( output['poly_class'].cpu().numpy(), data['mask'].cpu().numpy(), output['pred_polys'].cpu().numpy(), self.grid_size) # Get IoU iou = 0 pred_polys = output['pred_polys'].cpu().numpy() gt_polys = data['full_poly'] for i in range(pred_polys.shape[0]): p = pred_polys[i] p = utils.get_masked_poly(p, self.grid_size) p = utils.class_to_xy(p, self.grid_size) i, masks = metrics.iou_from_poly(p, gt_polys[i], self.grid_size, self.grid_size) iou += i iou = iou / pred_polys.shape[0] ious.append(iou) accuracies.append(accuracy) del (output) iou = np.mean(ious) accuracy = np.mean(accuracies) self.val_writer.add_scalar('iou', float(iou), self.global_step) self.val_writer.add_scalar('accuracy', float(accuracy), self.global_step) print '[VAL] IoU: %f, Accuracy: %f' % (iou, accuracy) # Reset self.model.train()
def validate(self): print 'Validating' self.model.eval() # Leave LSTM in train mode with torch.no_grad(): ious = [] for step, data in enumerate(tqdm(self.val_loader)): if len(data['orig_poly']) == 1: continue if self.opts['get_point_annotation']: img = data['img'].to(device) annotation = data['annotation_prior'].to(device).unsqueeze( 1) img = torch.cat([img, annotation], 1) else: img = data['img'].to(device) output = self.model.forward(img, data['fwd_poly']) pred_cps = output['pred_polys'][-1] pred_polys = self.spline.sample_point(pred_cps) pred_polys = pred_polys.data.cpu().numpy() # print(pred_polys.shape) # Get IoU iou = 0 orig_poly = data['orig_poly'] for i in range(pred_polys.shape[0]): curr_pred_poly = utils.poly01_to_poly0g( pred_polys[i], self.model.grid_size) curr_gt_poly = utils.poly01_to_poly0g( orig_poly[i], self.model.grid_size) i, masks = metrics.iou_from_poly( np.array(curr_pred_poly, dtype=np.int32), np.array(curr_gt_poly, dtype=np.int32), self.model.grid_size, self.model.grid_size) iou += i iou = iou / pred_polys.shape[0] ious.append(iou) del (output) del (pred_polys) iou = np.mean(ious) self.val_writer.add_scalar('iou', float(iou), self.global_step) print '[VAL] IoU: %f' % iou masks = np.tile(masks, [1, 1, 1, 3]) # Make [2, H, W, 3] img = (data['img'].cpu().numpy()[-1, ...] * 255).astype(np.uint8) img = np.transpose(img, [1, 2, 0]) # Make [H, W, 3] #self.val_writer.add_image('pred_mask', masks[0], self.global_step) #self.val_writer.add_image('gt_mask', masks[1], self.global_step) #self.val_writer.add_image('image', img, self.global_step) self.model.train()
def train(self, epoch): print 'Starting training' self.model.train() accum = defaultdict(float) # To accumulate stats for printing for step, data in enumerate(self.train_loader): if len(data['img']) == 1: continue if self.opts['get_point_annotation']: img = data['img'].to(device) annotation = data['annotation_prior'].to(device).unsqueeze(1) img = torch.cat([img, annotation], 1) else: img = data['img'].to(device) self.optimizer.zero_grad() if self.global_step % self.opts['val_freq'] == 0 and not self.opts[ 'debug']: #self.validate() self.save_checkpoint(epoch) output = self.model(img, data['fwd_poly'], data['sampled_interactive']) output_prob = F.softmax(output['x_prob'], dim=1) loss_sum = 0 pred_cps = output['pred_polys'][-1] pred_polys = self.spline.sample_point(pred_cps) #(bs, p_num, 2) #pred_polys = pred_cps # loss 1 gt_right_order, poly_mathcing_loss_sum = losses.poly_mathcing_loss( self.opts['p_num'], pred_polys, data['gt_poly'].to(device), loss_type=self.opts['loss_type']) loss_sum += poly_mathcing_loss_sum edge_annotation_loss = 0 # loss 2 curr_fp_edge_loss = self.opts['fp_weight'] * losses.fp_edge_loss( data['edge_mask'].to(device), output['edge_logits']) edge_annotation_loss += curr_fp_edge_loss # loss 3 fp_vertex_loss = self.opts['fp_weight'] * losses.fp_vertex_loss( data['vertex_mask'].to(device), output['vertex_logits']) # loss 4 , Dice loss, Boundary loss # dice_loss = losses.GeneralizedDice(output_prob, data['onehot_label'].cuda()) boundary_loss = losses.SurfaceLoss(output_prob, data['mask_distmap'].cuda()) edge_annotation_loss += fp_vertex_loss loss_sum += edge_annotation_loss #loss_sum += dice_loss #loss_sum += boundary_loss loss_sum = loss_sum + boundary_loss loss_sum.backward() if 'grad_clip' in self.opts.keys(): nn.utils.clip_grad_norm_(self.model.parameters(), self.opts['grad_clip']) self.optimizer.step() preds = pred_polys.detach().data.cpu().numpy() with torch.no_grad(): # Get IoU iou = 0 orig_poly = data['orig_poly'] for i in range(preds.shape[0]): curr_pred_poly = np.floor(preds[i] * 224).astype(np.int32) curr_gt_poly = np.floor(orig_poly[i] * 224).astype( np.int32) cur_iou, masks = metrics.iou_from_poly( np.array(curr_pred_poly, dtype=np.int32), np.array(curr_gt_poly, dtype=np.int32), 224, 224) iou += cur_iou iou = iou / preds.shape[0] accum['loss'] += float(loss_sum.item()) accum['iou'] += iou #accum['dice'] += float(dice_loss) accum['boundary'] += float(boundary_loss) accum['length'] += 1 if self.opts['edge_loss']: accum['edge_annotation_loss'] += float( edge_annotation_loss.item()) print( "[%s] Epoch: %d, Step: %d, Polygon Loss: %f, IOU: %f, Boundary: %f" \ % (str(datetime.now()), epoch, self.global_step, accum['loss'] / accum['length'], accum['iou'] / accum['length'], accum['boundary'])) if step % self.opts['print_freq'] == 0: # Mean of accumulated values for k in accum.keys(): if k == 'length': continue accum[k] /= accum['length'] # Add summaries masks = np.expand_dims(masks, -1).astype( np.uint8) # Add a channel dimension masks = np.tile(masks, [1, 1, 1, 3]) # Make [2, H, W, 3] img = (data['img'].cpu().numpy()[-1, ...] * 255).astype( np.uint8) img = np.transpose(img, [1, 2, 0]) # Make [H, W, 3] self.writer.add_image('pred_mask', masks[0], self.global_step) self.writer.add_image('gt_mask', masks[1], self.global_step) self.writer.add_image('image', img, self.global_step) for k in accum.keys(): if k == 'length': continue self.writer.add_scalar(k, accum[k], self.global_step) print( "[%s] Epoch: %d, Step: %d, Polygon Loss: %f, IOU: %f" \ % (str(datetime.now()), epoch, self.global_step, accum['loss'], accum['iou'])) accum = defaultdict(float) del (output, masks, pred_polys, preds, loss_sum) self.global_step += 1
def train(self, epoch): print 'Starting training' self.model.train() accum = defaultdict(float) #loss_loc = WHD_losses.WeightedHausdorffDistance(resized_height=self.opts['p_num'],resized_width=self.opts['p_num'], # return_2_terms=True, # device=device) #loss_loc = WHD_losses.AveragedHausdorffLoss() loss_loc = losses.WeightedHausdorffDistance(resized_height=224, resized_width=224, return_2_terms=True, device=device) # focalloss = focal.FocalLoss(None,None,None,'mean') focalloss = focal.FocalLoss() # To accumulate stats for printing for step, data in enumerate(self.train_loader): if len(data['img']) == 1: continue if self.opts['get_point_annotation']: img = data['img'].to(device) annotation = data['annotation_prior'].to(device).unsqueeze(1) img = torch.cat([img, annotation], 1) else: img = data['img'].to(device) self.optimizer.zero_grad() if self.global_step % self.opts['val_freq'] == 0 and not self.opts[ 'debug']: self.validate() self.save_checkpoint(epoch) output = self.model.forward(img, data['fwd_poly']) loss_sum = 0 pred_cps = output['pred_polys'][-1] pred_polys = self.spline.sample_point(pred_cps) # print(pred_polys.shape) # print(output['vertex_logits'].shape) gt_right_order, poly_mathcing_loss_sum = losses.poly_mathcing_loss( self.opts['p_num'], pred_polys, data['gt_poly'].to(device), loss_type=self.opts['loss_type']) # add by dzh contour refine ## Initializing Contour Box level_set_config_dict = { 'step_ckpts': [50], 'lambda_': 0.0, 'alpha': 1, 'smoothing': 1, 'render_radius': -1, 'is_gt_semantic': True, 'method': 'MLS', 'balloon': 1, 'threshold': 0.99, 'merge_weight': 0.5 } cbox = ContourBox.LevelSetAlignment(n_workers=1, fn_post_process_callback=None, config=level_set_config_dict) # print('-------------shape--------------------') output_contour, _ = cbox( { 'seg': np.expand_dims(data['edge_mask'], 0), 'bdry': None }, np.expand_dims( output['edge_logits'].view( data['edge_mask'].shape).cpu().detach().numpy(), 0)) masks_step = output_contour[0, :, 0, :, :] #--------add by dzh 7.18 edge_annotation_loss = 0 curr_fp_edge_loss = losses.fp_edge_loss( torch.from_numpy(masks_step).to( device ), #self.opts['fp_weight'] * losses.fp_edge_loss(torch.from_numpy(masks_step).to(device), output['edge_logits'] ) #data['edge_mask'] torch.from_numpy(masks_step) edge_annotation_loss += curr_fp_edge_loss tt = [] #pred_poly_mask = np.zeros((36, 36), np.float32) for i in range(pred_polys.shape[0]): pred_poly_mask = np.zeros((224, 224), dtype=np.float32) ff = np.floor(pred_polys[i].detach().cpu().numpy() * 36).astype(np.int32) if not isinstance(ff, list): ff = [ff] for p in ff: pred_poly_mask = utils.draw_poly(pred_poly_mask, p) #ff=utils.poly01_to_poly0g(pred_polys[i].detach().cpu().numpy(), 35) # pred_poly_mask = utils.get_vertices_mask_36(ff, pred_poly_mask) tt.append(pred_poly_mask) tt1 = np.array(tt, dtype=np.float32) pred_poly_mask11 = torch.from_numpy(tt1).cuda() ll1 = pred_poly_mask11 #ll1 = output['vertex_logits'].view(output['vertex_logits'].shape[0],28,28) jjj = [] for tt in range(ll1.shape[0]): jjj.append([224, 224]) #jjj = [[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28],[28,28]] # print(data['poly_mask'].shape) kk = [] poly_mask_ori = data['poly_mask'] for hh in range(ll1.shape[0]): zzz = torch.FloatTensor(poly_mask_ori[hh].astype( np.float32)).cuda() #zzz = torch.FloatTensor(data['gt_orig_poly'][hh]).cuda() kk.append(zzz) #zzz = torch.from_numpy(data['gt_orig_poly'][0]) # print(ll1.shape) # print(kk.shape) #ll1,kk term1, term2 = loss_loc.forward( ll1, kk, torch.FloatTensor(np.array(jjj, dtype=np.float32)).cuda()) #fp_vertex_loss = self.opts['fp_weight'] * (term1+term2) fp_vertex_loss = 0.1 * (term1 + term2) + poly_mathcing_loss_sum #fp_vertex_loss = poly_mathcing_loss_sum + self.opts['fp_weight']* 0.1 * (term1+term2) loss_sum += fp_vertex_loss loss_sum += edge_annotation_loss # + self.opts['fp_weight'] * (term1+term2) ################iou loss function################# #preds= pred_polys.detach().data.cpu().numpy() #iou_loss = 0 #orig_poly = data['orig_poly'] #for i in range(preds.shape[0]): # curr_pred_poly = np.floor(preds[i] * 224).astype(np.int32) # curr_gt_poly = np.floor(orig_poly[i] * 224).astype(np.int32) # cur_iou, masks = metrics.iou_from_poly(np.array(curr_pred_poly, dtype=np.int32), # np.array(curr_gt_poly, dtype=np.int32), # 224, 224) # iou_loss += cur_iou #iou_loss = -iou_loss / preds.shape[0] #loss_sum += 0.1 * iou_loss ################iou loss function################# with torch.no_grad(): iou = 0 gt_mask_0 = [] pred_mask_0 = [] orig_poly = data['orig_poly'] preds = pred_polys.detach().data.cpu().numpy() # iou_filter = [] for i in range(preds.shape[0]): curr_pred_poly = np.floor(preds[i] * 224).astype(np.int32) curr_gt_poly = np.floor(orig_poly[i] * 224).astype( np.int32) cur_iou, masks = metrics.iou_from_poly( np.array(curr_pred_poly, dtype=np.int32), np.array(curr_gt_poly, dtype=np.int32), 224, 224) gt_mask_0.append(masks[1]) pred_mask_0.append(masks[0]) gt_mask_1 = torch.from_numpy( np.array(gt_mask_0)).to(device).float() pred_mask_1 = torch.from_numpy( np.array(pred_mask_0)).to(device).float() # mask_loss = focalloss(pred_mask_1, gt_mask_1) # mask_loss = losses.class_balanced_cross_entropy_loss(pred_mask_1, gt_mask_1) # pred111=pred_mask_1.view(pred_mask_1.shape[0],1,224,224) #mask_loss = 100 * focalloss((pred_mask_1/255), (gt_mask_1/255)) mask_loss = torch.sum( torch.abs(gt_mask_1 / 250 - pred_mask_1 / 250)) loss_sum += torch.mean(mask_loss) # # iou_filter.append(1 if cur_iou>self.opts['iou_filter'] else 0) # iou += cur_iou # iou = iou / preds.shape[0] # # iou_filter = np.array(iou_filter) # # iou_filter = torch.from_numpy(iou_filter).to(device).float() # loss_sum += (-iou) # if self.opts['iou_filter']>0: # loss_sum = (loss_sum + (1-iou)) * iou_filter # loss_sum = torch.mean(loss_sum) loss_sum.backward() if 'grad_clip' in self.opts.keys(): nn.utils.clip_grad_norm_(self.model.parameters(), self.opts['grad_clip']) self.optimizer.step() preds = pred_polys.detach().data.cpu().numpy() with torch.no_grad(): # Get IoU iou = 0 orig_poly = data['orig_poly'] for i in range(preds.shape[0]): curr_pred_poly = np.floor(preds[i] * 224).astype(np.int32) curr_gt_poly = np.floor(orig_poly[i] * 224).astype( np.int32) cur_iou, masks = metrics.iou_from_poly( np.array(curr_pred_poly, dtype=np.int32), np.array(curr_gt_poly, dtype=np.int32), 224, 224) iou += cur_iou iou = iou / preds.shape[0] accum['loss'] += float(loss_sum.item()) accum['iou'] += iou accum['length'] += 1 if self.opts['edge_loss']: accum['edge_annotation_loss'] += float( edge_annotation_loss.item()) print( "[%s] Epoch: %d, Step: %d, Polygon Loss: %f, IOU: %f" \ % (str(datetime.now()), epoch, self.global_step, accum['loss'] / accum['length'], accum['iou'] / accum['length'])) if step % self.opts['print_freq'] == 0: # Mean of accumulated values for k in accum.keys(): if k == 'length': continue accum[k] /= accum['length'] # Add summaries masks = np.expand_dims(masks, -1).astype( np.uint8) # Add a channel dimension #print(masks.shape) masks = np.tile(masks, [1, 1, 1, 3]) # Make [2, H, W, 3] img = (data['img'].cpu().numpy()[-1, ...] * 255).astype( np.uint8) img = np.transpose(img, [1, 2, 0]) # Make [H, W, 3] self.writer.add_image('pred_mask', masks[0], self.global_step) self.writer.add_image('gt_mask', masks[1], self.global_step) self.writer.add_image('image', img, self.global_step) self.writer.add_image( 'edge_acm_gt', np.tile( np.expand_dims(masks_step[preds.shape[0] - 1], axis=-1).astype(np.uint8), [1, 1, 3]), self.global_step) #self.writer.add_image('ori_GT', pred_edge_mask = np.tile( np.expand_dims( output['edge_logits'].cpu().numpy()[preds.shape[0] - 1] * 255, axis=-1).astype(np.uint8), [1, 1, 3]).reshape(28, 28, 3) #print(pred_edge_mask.shape) self.writer.add_image('pred_edge', pred_edge_mask, self.global_step) for k in accum.keys(): if k == 'length': continue self.writer.add_scalar(k, accum[k], self.global_step) print( "[%s] Epoch: %d, Step: %d, Polygon Loss: %f, IOU: %f" \ % (str(datetime.now()), epoch, self.global_step, accum['loss'], accum['iou'])) accum = defaultdict(float) del (output, masks, pred_polys, preds, loss_sum) self.global_step += 1
def train(self, epoch): print('Starting training') self.model.train() accum = defaultdict(float) # To accumulate stats for printin for step, data in enumerate(self.train_loader): if self.global_step % self.opts['val_freq'] == 0: self.validate() self.save_checkpoint(epoch) # Forward pass output = self.model(data['img'].to(device), data['fwd_poly'].to(device)) # Smoothed targets dt_targets = utils.dt_targets_from_class(output['poly_class'].cpu().numpy(), self.grid_size, self.opts['dt_threshold']) # Get losses loss = losses.poly_vertex_loss_mle(torch.from_numpy(dt_targets).to(device), data['mask'].to(device), output['logits']) fp_edge_loss = self.opts['fp_weight'] * losses.fp_edge_loss(data['edge_mask'].to(device), output['edge_logits']) fp_vertex_loss = self.opts['fp_weight'] * losses.fp_vertex_loss(data['vertex_mask'].to(device), output['vertex_logits']) total_loss = loss + fp_edge_loss + fp_vertex_loss # Backward pass self.optimizer.zero_grad() total_loss.backward() if 'grad_clip' in self.opts.keys(): nn.utils.clip_grad_norm_(self.model.parameters(), self.opts['grad_clip']) self.optimizer.step() # Get accuracy accuracy = metrics.train_accuracy(output['poly_class'].cpu().numpy(), data['mask'].cpu().numpy(), output['pred_polys'].cpu().numpy(), self.grid_size) # Get IoU iou = 0 pred_polys = output['pred_polys'].cpu().numpy() gt_polys = data['full_poly'] for i in range(pred_polys.shape[0]): p = pred_polys[i] p = utils.get_masked_poly(p, self.grid_size) p = utils.class_to_xy(p, self.grid_size) i, masks = metrics.iou_from_poly(p, gt_polys[i], self.grid_size, self.grid_size) iou += i iou = iou / pred_polys.shape[0] accum['loss'] += float(loss) accum['fp_edge_loss'] += float(fp_edge_loss) accum['fp_vertex_loss'] += float(fp_vertex_loss) accum['accuracy'] += accuracy accum['iou'] += iou accum['length'] += 1 if step % self.opts['print_freq'] == 0: # Mean of accumulated values for k in accum.keys(): if k == 'length': continue accum[k] /= accum['length'] # Add summaries masks = np.expand_dims(masks, -1).astype(np.uint8) # Add a channel dimension masks = np.tile(masks, [1, 1, 1, 3]) # Make [2, H, W, 3] img = (data['img'].cpu().numpy()[-1,...]*255).astype(np.uint8) img = np.transpose(img, [1,2,0]) # Make [H, W, 3] vert_logits = np.reshape(output['vertex_logits'][-1, ...].detach().cpu().numpy(), (self.grid_size, self.grid_size, 1)) edge_logits = np.reshape(output['edge_logits'][-1, ...].detach().cpu().numpy(), (self.grid_size, self.grid_size, 1)) vert_logits = (1/(1 + np.exp(-vert_logits))*255).astype(np.uint8) edge_logits = (1/(1 + np.exp(-edge_logits))*255).astype(np.uint8) vert_logits = np.tile(vert_logits, [1, 1, 3]) # Make [H, W, 3] edge_logits = np.tile(edge_logits, [1, 1, 3]) # Make [H, W, 3] vertex_mask = np.tile(np.expand_dims(data['vertex_mask'][-1,...].cpu().numpy().astype(np.uint8)*255,-1),(1,1,3)) edge_mask = np.tile(np.expand_dims(data['edge_mask'][-1,...].cpu().numpy().astype(np.uint8)*255,-1),(1,1,3)) self.writer.add_image('pred_mask', masks[0], self.global_step) self.writer.add_image('gt_mask', masks[1], self.global_step) self.writer.add_image('image', img, self.global_step) self.writer.add_image('vertex_logits', vert_logits, self.global_step) self.writer.add_image('edge_logits', edge_logits, self.global_step) self.writer.add_image('edge_mask', edge_mask, self.global_step) self.writer.add_image('vertex_mask', vertex_mask, self.global_step) if self.opts['return_attention'] is True: att = output['attention'][-1, 1:4, ...].detach().cpu().numpy() att = np.transpose(att, [0, 2, 3, 1]) # Make [T, H, W, 1] att = np.tile(att, [1, 1, 1, 3]) # Make [T, H, W, 3] def _scale(att): att = att/np.max(att) return (att*255).astype(np.int32) self.writer.add_image('attention_1', pyramid_expand(_scale(att[0]), upscale=8, sigma=10), self.global_step) self.writer.add_image('attention_2', pyramid_expand(_scale(att[1]), upscale=8, sigma=10), self.global_step) self.writer.add_image('attention_3', pyramid_expand(_scale(att[2]), upscale=8, sigma=10), self.global_step) for k in accum.keys(): if k == 'length': continue self.writer.add_scalar(k, accum[k], self.global_step) print("[%s] Epoch: %d, Step: %d, Polygon Loss: %f, Edge Loss: %f, Vertex Loss: %f, Accuracy: %f, IOU: %f"\ %(str(datetime.now()), epoch, self.global_step, accum['loss'], accum['fp_edge_loss'], accum['fp_vertex_loss'],\ accum['accuracy'], accum['iou'])) accum = defaultdict(float) del(output) self.global_step += 1
def train(self, epoch): print 'Starting training' self.model.temperature = self.opts['temperature'] self.model.ggnn.encoder.train() accum = defaultdict(float) # To accumulate stats for printing ggnn_grid_size = self.opts['ggnn_grid_size'] for step, data in enumerate(self.train_loader): self.optimizer.zero_grad() if self.global_step % self.opts['val_freq'] == 0: self.validate() self.save_checkpoint(epoch) output = self.model(data['img'].to(device), data['fwd_poly'].to(device), orig_poly=data['orig_poly']) ggnn_logits = output['ggnn_logits'] local_prediction = output['ggnn_local_prediction'].to(device) poly_masks = output['ggnn_mask'].to(device) pred_polys = output['pred_polys'].data.numpy() loss_sum = losses.poly_vertex_loss_mle_ggnn( local_prediction, poly_masks, ggnn_logits) loss_sum.backward() if 'grad_clip' in self.opts.keys(): # "grad_clip": 40 nn.utils.clip_grad_norm_(self.model.ggnn.parameters(), self.opts['grad_clip']) self.optimizer.step() with torch.no_grad(): # Get IoU iou = 0 orig_poly = data['orig_poly'] for i in range(pred_polys.shape[0]): p = pred_polys[i] mask_poly = utils.get_masked_poly( p, self.model.ggnn.ggnn_grid_size) #"ggnn_grid_size": 112 mask_poly = utils.class_to_xy( mask_poly, self.model.ggnn.ggnn_grid_size) curr_gt_poly_112 = utils.poly01_to_poly0g( orig_poly[i], ggnn_grid_size) cur_iou, masks = metrics.iou_from_poly( np.array(mask_poly, dtype=np.int32), np.array(curr_gt_poly_112, dtype=np.int32), ggnn_grid_size, ggnn_grid_size) iou += cur_iou iou = iou / pred_polys.shape[0] accum['loss'] += float(loss_sum.item()) accum['iou'] += iou accum['length'] += 1 if step % self.opts['print_freq'] == 0: #"print_freq": 20 # Mean of accumulated values for k in accum.keys(): if k == 'length': continue accum[k] /= accum['length'] # Add summaries masks = np.expand_dims(masks, -1).astype( np.uint8) # Add a channel dimension masks = np.tile(masks, [1, 1, 1, 3]) # Make [2, H, W, 3] img = (data['img'].cpu().numpy()[-1, ...] * 255).astype( np.uint8) img = np.transpose( img, [1, 2, 0]) # Make [H, W, 3], swap the dimention self.writer.add_image('pred_mask', masks[0], self.global_step) self.writer.add_image('gt_mask', masks[1], self.global_step) self.writer.add_image('image', img, self.global_step) for k in accum.keys(): if k == 'length': continue self.writer.add_scalar(k, accum[k], self.global_step) print( "[%s] Epoch: %d, Step: %d, Polygon Loss: %f, IOU: %f" \ % (str(datetime.now()), epoch, self.global_step, accum['loss'], accum['iou'])) accum = defaultdict(float) del (output, local_prediction, poly_masks, masks, ggnn_logits, pred_polys, loss_sum) self.global_step += 1