def test(val_loader, model, epoch): model.eval() step = 0 batch_time = AverageMeter() total_losses = AverageMeter() binary_losses = AverageMeter() instance_losses = AverageMeter() mean_iou = AverageMeter() end = time.time() val_img_list = [] # val_img_md5 = open(os.path.join(im_path, "val_" + str(epoch + 1) + ".txt"), "w") for batch_idx, batch in enumerate(iter(val_loader)): step += 1 # image_data = Variable(input_data["input_tensor"]).to(DEVICE) # instance_label = Variable(input_data["instance_label"]).to(DEVICE) # binary_label = Variable(input_data["binary_label"]).to(DEVICE) image_data = Variable(batch[0]).type(torch.FloatTensor).to(DEVICE) binary_label = Variable(batch[1]).type(torch.LongTensor).to(DEVICE) instance_label = Variable(batch[2]).type(torch.FloatTensor).to(DEVICE) # output process net_output = model(image_data) net_output total_loss, binary_loss, instance_loss, out, val_iou = compute_loss( net_output, binary_label, instance_label) total_losses.update(total_loss.item(), image_data.size()[0]) binary_losses.update(binary_loss.item(), image_data.size()[0]) instance_losses.update(instance_loss.item(), image_data.size()[0]) mean_iou.update(val_iou, image_data.size()[0]) # if step % 100 == 0: # val_img_list.append( # compose_img(image_data, out, binary_label, net_output["instance_seg_logits"], instance_label, 0)) # val_img_md5.write(input_data["img_name"][0] + "\n") # lane_cluster_and_draw(image_data, net_output["binary_seg_pred"], net_output["instance_seg_logits"], input_data["o_size"], input_data["img_name"], json_path) batch_time.update(time.time() - end) end = time.time() # print( # "Epoch {ep} Validation Report | ETA: {et:.2f}|Total:{tot:.5f}|Binary:{bin:.5f}|Instance:{ins:.5f}|IoU:{iou:.5f}".format( # ep=epoch + 1, # et=batch_time.val, # tot=total_losses.avg, # bin=binary_losses.avg, # ins=instance_losses.avg, # iou=mean_iou.avg, # )) # sys.stdout.flush() # val_img = np.concatenate(val_img_list, axis=1) # cv2.imwrite(os.path.join(im_path, "val_" + str(epoch + 1) + ".png"), val_img) # val_img_md5.close() return mean_iou.avg
def train(train_loader, model, optimizer, epoch, w1, w2, w3, w4): model.train() batch_time = AverageMeter() mean_iou = AverageMeter() total_losses = AverageMeter() binary_losses = AverageMeter() instance_losses = AverageMeter() end = time.time() step = 0 t = tqdm(enumerate(iter(train_loader)), leave=False, total=len(train_loader)) for batch_idx, batch in t: try: step += 1 image_data = Variable(batch[0]).type(torch.FloatTensor).to(DEVICE) binary_label = Variable(batch[1]).type(torch.LongTensor).to(DEVICE) instance_label = Variable(batch[2]).type( torch.FloatTensor).to(DEVICE) #print("///////////////////////////////////////////////////") #print(image_data.size()) #print(binary_label.size()) # # print(image_data.shape) #print("///////////////////////////////////////////////////") # forward pass net_output = model(image_data) # compute loss total_loss, binary_loss, instance_loss, out, train_iou = compute_loss( net_output, binary_label, instance_label, w1, w2, w3, w4) # update loss in AverageMeter instance total_losses.update(total_loss.item(), image_data.size()[0]) binary_losses.update(binary_loss.item(), image_data.size()[0]) instance_losses.update(instance_loss.item(), image_data.size()[0]) mean_iou.update(train_iou, image_data.size()[0]) # reset gradients optimizer.zero_grad() # backpropagate total_loss.backward() # update weights optimizer.step() # update batch time batch_time.update(time.time() - end) end = time.time() if step % config.show_interval == 0: print( "Epoch {ep} Step {st} |({batch}/{size})| ETA: {et:.2f}|Total loss:{tot:.5f}|Binary loss:{bin:.5f}|Instance loss:{ins:.5f}|IoU:{iou:.5f}" .format( ep=epoch + 1, st=step, batch=batch_idx + 1, size=len(train_loader), et=batch_time.val, tot=total_losses.avg, bin=binary_losses.avg, ins=instance_losses.avg, iou=train_iou, )) print("current learning rate is %s" % (str(optimizer.state_dict()['param_groups'][0]['lr']))) sys.stdout.flush() train_img_list = [] for i in range(3): train_img_list.append( compose_img(image_data, out, binary_label, net_output["instance_seg_logits"], instance_label, i)) train_img = np.concatenate(train_img_list, axis=1) cv2.imwrite( os.path.join( "./output", "train_" + str(epoch + 1) + "_step_" + str(step) + ".png"), train_img) except Exception as e: print(e) print('error') return mean_iou.avg
def train(train_loader, model, optimizer, epoch, image_output_path): batch_time = AverageMeter() mean_iou = AverageMeter() total_losses = AverageMeter() binary_losses = AverageMeter() instance_losses = AverageMeter() end = time.time() step = 0 t = tqdm(enumerate(iter(train_loader)), leave=False, total=len(train_loader)) for batch_idx, batch in t: # MES changes to use less GPU memory step += 1 image_data = Variable(batch[0]).type(torch.FloatTensor).to(DEVICE) binary_label = Variable(batch[1]).type(torch.LongTensor).to(DEVICE) instance_label = Variable(batch[2]).type(torch.FloatTensor).to(DEVICE) # forward pass net_output = model(image_data) # compute loss total_loss, binary_loss, instance_loss, out, train_iou = compute_loss( net_output, binary_label, instance_label) # update loss in AverageMeter instance total_losses.update(total_loss.item(), image_data.size()[0]) binary_losses.update(binary_loss.item(), image_data.size()[0]) instance_losses.update(instance_loss.item(), image_data.size()[0]) mean_iou.update(train_iou, image_data.size()[0]) # reset gradients optimizer.zero_grad() # backpropagate total_loss.backward() # update weights optimizer.step() # update batch time batch_time.update(time.time() - end) end = time.time() if step % 1 == 0: print( "Epoch {ep} Step {st} |({batch}/{size})| ETA: {et:.2f}|Total loss:{tot:.5f}|Binary loss:{bin:.5f}|Instance loss:{ins:.5f}|IoU:{iou:.5f}" .format( ep=epoch + 1, st=step, batch=batch_idx + 1, size=len(train_loader), et=batch_time.val, tot=total_losses.avg, bin=binary_losses.avg, ins=instance_losses.avg, iou=train_iou, )) sys.stdout.flush() train_img_list = [] for i in range(3): train_img_list.append( compose_img(image_data, out, binary_label, net_output["instance_seg_logits"], instance_label, i)) train_img = np.concatenate(train_img_list, axis=1) # MES changes print("train: about to imwrite images") print("train_img: ", train_img) print(" ") cv2.imwrite( os.path.join( image_output_path, "train_" + str(epoch + 1) + "_step_" + str(step) + ".png"), train_img) return mean_iou.avg