def train(num_gpu, train_loader, model, criterion, optimizer, lovasz, epoch, total_ep): ''' :param train_loader: loaded for training dataset :param model: model :param criterion: loss function :param optimizer: optimization algo, such as ADAM or SGD :param epoch: epoch number :return: average epoch loss, overall pixel-wise accuracy, per class accuracy, per class iu, and mIOU ''' # switch to train mode model.train() mIOU = [] epoch_loss = [] total_time = 0 total_batches = len(train_loader) for i, (input, target) in enumerate(train_loader): start_time = time.time() if num_gpu > 0: input = input.cuda() target = target.cuda() input_var = torch.autograd.Variable(input) target_var = torch.autograd.Variable(target) # run the mdoel output = model(input_var) # set the grad to zero optimizer.zero_grad() loss = criterion(output, target_var) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss.append(loss.item()) time_taken = time.time() - start_time total_time += time_taken # compute the confusion matrix if lovasz: IOU = iou_binary((output.data > 0).long(), target_var) else: IOU = iou_binary(output.max(1)[1].data.long(), target_var) mIOU.append(IOU) average_epoch_loss_train = sum(epoch_loss) / len(epoch_loss) print('[%d/%d] loss: %.3f time:%.2f' % (epoch + 1, total_ep, average_epoch_loss_train, total_time / total_batches)) return average_epoch_loss_train, sum(mIOU) / len(mIOU),
def train(num_gpu, train_loader, model, criterion, optimizer, lovasz, epoch, total_ep): ''' :param train_loader: loaded for training dataset :param model: model :param criterion: loss function :param optimizer: optimization algo, such as ADAM or SGD :param epoch: epoch number :return: average epoch loss, overall pixel-wise accuracy, per class accuracy, per class iu, and mIOU ''' # switch to train mode model.train() mIOU = [] epoch_loss = [] total_time = 0 total_batches = len(train_loader) for i, (input, target) in enumerate(train_loader): # input_arr = (6500*input[0].data.permute(1, 2, 0).numpy()+128).astype(np.uint8) # cv2.imshow("image0", input_arr) # target_arr = 255 * target[0].data.numpy().astype(np.uint8) # cv2.imshow("label0", target_arr) # # input_arr = (6500 * input[1].data.permute(1, 2, 0).numpy() + 128).astype(np.uint8) # cv2.imshow("image1", input_arr) # target_arr = 255 * target[1].data.numpy().astype(np.uint8) # cv2.imshow("label1", target_arr) # # input_arr = (6500 * input[2].data.permute(1, 2, 0).numpy() + 128).astype(np.uint8) # cv2.imshow("image2", input_arr) # target_arr = 255 * target[2].data.numpy().astype(np.uint8) # cv2.imshow("label2", target_arr) # # input_arr = (6500 * input[3].data.permute(1, 2, 0).numpy() + 128).astype(np.uint8) # cv2.imshow("image3", input_arr) # target_arr = 255 * target[3].data.numpy().astype(np.uint8) # cv2.imshow("label3", target_arr) # # cv2.waitKey(0) # cv2.destroyAllWindows() # # start_time = time.time() # input_arr[:,:,0] = input_arr[:,:,0] # input_arr = torchvision.transforms.ToPILImage()(input[0]) # cv2.imshow("image", np.array(input_arr)) # [106.90267 115.0241 130.67285] # [62.213623 64.17969 68.387596] if num_gpu > 0: input = input.cuda() target = target.cuda() input_var = torch.autograd.Variable(input) target_var = torch.autograd.Variable(target) # run the mdoel output = model(input_var) # set the grad to zero optimizer.zero_grad() # print(input.size()) # print(output.size()) # print(target_var.size()) loss = criterion(output, target_var) optimizer.zero_grad() loss.backward() optimizer.step() if int(torch.__version__[2]) < 4: epoch_loss.append(loss.data[0]) else: epoch_loss.append(loss.item()) time_taken = time.time() - start_time total_time += time_taken # compute the confusion matrix if lovasz: IOU = iou_binary((output.data > 0).long(), target_var) else: IOU = iou_binary(output.max(1)[1].data.long(), target_var) mIOU.append(IOU) average_epoch_loss_train = sum(epoch_loss) / len(epoch_loss) print('[%d/%d] loss: %.3f time:%.2f' % (epoch + 1, total_ep, average_epoch_loss_train, total_time / total_batches)) return average_epoch_loss_train, sum(mIOU) / len(mIOU),
def val(num_gpu, val_loader, model, criterion, lovasz, vis=False): ''' :param val_loader: loaded for validation dataset :param model: model :param criterion: loss function :return: average epoch loss, overall pixel-wise accuracy, per class accuracy, per class iu, and mIOU ''' # switch to evaluation mode model.eval() total_time = 0 epoch_loss = [] total_batches = len(val_loader) rand_pick = np.random.randint(0, len(val_loader)) mIOU = [] f1score = [] for i, (input, target, edge_target) in enumerate(val_loader): start_time = time.time() if num_gpu > 0: input = input.cuda() target = target.cuda() edge_target = edge_target.cuda() with torch.no_grad(): input_var, target_var, edge_target_var \ = torch.autograd.Variable(input), torch.autograd.Variable(target), \ torch.autograd.Variable(edge_target) # run the mdoel output = model(input_var) # compute the loss loss = criterion(output, target_var) if int(torch.__version__[2]) < 4: epoch_loss.append(loss.data[0]) else: epoch_loss.append(loss.item()) time_taken = time.time() - start_time total_time += time_taken # compute the confusion matrix if lovasz: IOU = iou_binary((output.data > 0).long(), target_var) f1 = calcF1((output.data > 0).long(), edge_target_var, ignore=255, per_image=True) else: IOU = iou_binary(output.max(1)[1].data.long(), target_var) f1 = calcF1(output.max(1)[1].data.long(), edge_target_var, ignore=255, per_image=True) mIOU.append(IOU) f1score.append(f1) if i == rand_pick and vis: save_input = input_var save_out = output save_target = target_var average_epoch_loss_val = sum(epoch_loss) / len(epoch_loss) print('loss: %.3f time:%.2f f1score: %.2f ' % (average_epoch_loss_val, total_time / total_batches, sum(f1score) / len(f1score))) if vis: return average_epoch_loss_val, sum(mIOU) / len( mIOU), save_input, save_out, save_target else: return average_epoch_loss_val, sum(mIOU) / len(mIOU)