# -*- coding: utf-8 -*- """ttrain.ipynb Automatically generated by Colaboratory. Original file is located at https://colab.research.google.com/drive/1uwgolSLC-Ubd0o3PPxPFB7LnZhXYmQxr """ ! pip install einops ! pip install timm ! pip install tensorboardX load=None lr= 1e-4 import os import torch import torch.nn.functional as F import sys import numpy as np from datetime import datetime from torchvision.utils import make_grid from my_model_19 import IR_Net from data_edge import get_loader,test_dataset from utils_new import clip_gradient, adjust_lr from tensorboardX import SummaryWriter import logging import torch.backends.cudnn as cudnn import ResNet from sa import SpatialAttention gpu_id=0 #set the device for training if gpu_id=='0': os.environ["CUDA_VISIBLE_DEVICES"] = "0" print('USE GPU 0') elif gpu_id=='1': os.environ["CUDA_VISIBLE_DEVICES"] = "1" print('USE GPU 1') cudnn.benchmark = True # Commented out IPython magic to ensure Python compatibility. import torch from torch import nn from torch import functional as F from torch import optim import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader from tqdm import tqdm import math import matplotlib.pyplot as plt import os from pytorch_iou import IOU # %matplotlib inline import os os.makedirs('train_file') os.makedirs('test_file') os.makedirs('test_in_train_file') save_path='/content/train_file/RGBD_for_train/depth_after_HHA/' if not os.path.exists(save_path): os.makedirs(save_path) save_path1='/content/test_in_train_file/test_in_train/depth_after_HHA/' if not os.path.exists(save_path1): os.makedirs(save_path1) edge_path_train='/content/train_file/RGBD_for_train/edge_maps/' if not os.path.exists(edge_path_train): os.makedirs(edge_path_train) edge_path_val='/content/test_in_train_file/test_in_train/edge_maps/' if not os.path.exists(edge_path_val): os.makedirs(edge_path_val) # Unzipping training edge maps ! unzip "/content/drive/MyDrive/Colab Notebooks/RGBD_SOD/Datasets/edge_training_Canny_fromGTsalmap.zip" -d "/content/train_file/RGBD_for_train/edge_maps/" # Unzipping training edge maps #! unzip "/content/drive/MyDrive/Colab Notebooks/image_retargeting/Dataset_fromBBS_Net/edge_training_Canny_fromGTsalmap.zip" -d "/content/train_file/RGBD_for_train/edge_maps/" # Unzipping training images ok ! unzip "/content/drive/MyDrive/Colab Notebooks/RGBD_SOD/Datasets/RGBD_for_train.zip" -d "/content/train_file/" # Unzipping training images new depth maps ok ! unzip "/content/drive/MyDrive/Colab Notebooks/RGBD_SOD/Datasets/train_depth_after_HHA.zip" -d "/content/train_file/RGBD_for_train/depth_after_HHA/" #unzipping validation images ok ! unzip "/content/drive/MyDrive/Colab Notebooks/RGBD_SOD/Datasets/test_in_train.zip" -d "/content/test_in_train_file/" # Unzipping validation images new depth maps ok ! unzip "/content/drive/MyDrive/Colab Notebooks/RGBD_SOD/Datasets/val_depth_after_HHA.zip" -d "/content/test_in_train_file/test_in_train/depth_after_HHA/" #unzipping test images ok ! unzip "/content/drive/MyDrive/Colab Notebooks/RGBD_SOD/Datasets/RGBD_for_test.zip" -d "/content/test_file/" model= IR_Net(in_channels = 3, cmt_channel = [46, 92, 184, 368], patch_channel = [46, 92, 184, 368], block_layer = [2, 2, 10, 2], R = 3.6, img_size = 224) print(model) model.cuda() params = model.parameters() optimizer = torch.optim.Adam(params, lr) for parameter in model.parameters(): parameter.requires_grad = False model.cuda() model.eval() dim=(224,224) save_path='/content/test_file/RGBD_for_test/NLPR/depth_HHA/' if not os.path.exists(save_path): os.makedirs(save_path) ! unzip "/content/drive/MyDrive/Colab Notebooks/RGBD_SOD/Datasets/depth_HHA_NLPR.zip" -d "/content/test_file/RGBD_for_test/NLPR/depth_HHA/" model.cuda() params = model.parameters() optimizer = torch.optim.Adam(params, lr) #set the path image_root ='/content/train_file/RGBD_for_train/RGB/' gt_root = '/content/train_file/RGBD_for_train/GT/' depth_root='/content/train_file/RGBD_for_train/depth_after_HHA/' edge_root= '/content/train_file/RGBD_for_train/edge_maps/' # validation data path test_image_root='/content/test_in_train_file/test_in_train/RGB/' test_gt_root='/content/test_in_train_file/test_in_train/GT/' test_depth_root='/content/test_in_train_file/test_in_train/depth_after_HHA/' #val_edge_root= '/content/test_in_train_file/test_in_train/edge_maps/' save_path_cpoints='/content/chk_points/' if not os.path.exists(save_path_cpoints): os.makedirs(save_path_cpoints) batchsize=5 trainsize= 224 epoch=125 clip=0.5 decay_rate=0.1 decay_epoch=50 #load data print('load data...') train_loader = get_loader(image_root, gt_root,depth_root, batchsize=batchsize, trainsize=trainsize) test_loader = test_dataset(test_image_root, test_gt_root,test_depth_root, trainsize) total_step = len(train_loader) print("Loading done") logging.basicConfig(filename=save_path_cpoints+'log.log',format='[%(asctime)s-%(filename)s-%(levelname)s:%(message)s]', level = logging.INFO,filemode='a',datefmt='%Y-%m-%d %I:%M:%S %p') logging.info("BBSNet-Train") logging.info("Config") logging.info('epoch:{};lr:{};batchsize:{};trainsize:{};clip:{};decay_rate:{};save_path:{};decay_epoch:{}'.format(epoch,lr,batchsize,trainsize,clip,decay_rate,save_path_cpoints,decay_epoch)) #set loss function #CE = torch.nn.BCEWithLogitsLoss().cuda() CE = torch.nn.BCEWithLogitsLoss() IOU = IOU(size_average = True) #SSIM= SSIM(size_average = True) #ECE = torch.nn.BCELoss() #CE_levels = torch.nn.BCELoss().cuda() import torch.nn.functional as F criterion= torch.nn.BCEWithLogitsLoss() ## train function with edge maps added, no side supervision step=0 writer = SummaryWriter(save_path_cpoints+'summary') best_mae=1 best_epoch=0 clip=0.5 #train function def train(train_loader, model, optimizer, epoch,save_path): global step model.cuda() model.train() sal_loss_all = 0 loss_all = 0 epoch_step = 0 try: for i, (images, gts, depths) in enumerate(train_loader, start=1): optimizer.zero_grad() images = images.cuda() gts = gts.cuda() depths=depths.cuda() s1, s1_sig = model(images,depths) sal_loss= CE(s1,gts) +IOU(s1_sig, gts) loss = sal_loss loss.backward() clip_gradient(optimizer, clip) optimizer.step() step+=1 epoch_step+=1 loss_all+=float(loss.data) if i % 100 == 0 or i == total_step or i==1: print('{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], Loss1: {:.4f} '. format(datetime.now(), epoch, epoch, i, total_step, loss.data)) logging.info('#TRAIN#:Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], Loss1: {:.4f} '. format( epoch, epoch, i, total_step, loss.data)) writer.add_scalar('Loss', loss.data, global_step=step) grid_image = make_grid(images[0].clone().cpu().data, 1, normalize=True) writer.add_image('RGB', grid_image, step) grid_image = make_grid(gts[0].clone().cpu().data, 1, normalize=True) writer.add_image('Ground_truth', grid_image, step) res=s1[0].clone() res = res.sigmoid().data.cpu().numpy().squeeze() res = (res - res.min()) / (res.max() - res.min() + 1e-8) writer.add_image('s1', torch.tensor(res), step,dataformats='HW') loss_all/=epoch_step logging.info('#TRAIN#:Epoch [{:03d}/{:03d}], Loss_AVG: {:.4f}'.format( epoch, epoch, loss_all)) writer.add_scalar('Loss-epoch', loss_all, global_step=epoch) #if (epoch) % 5 == 0: # torch.save(model.state_dict(), save_path+'IR_Net_epoch_{}.pth'.format(epoch)) except KeyboardInterrupt: #print('Keyboard Interrupt: save model and exit.') #if not os.path.exists(save_path): # os.makedirs(save_path) #torch.save(model.state_dict(), save_path+'IR_Net_epoch_{}.pth'.format(epoch+1)) #print('save checkpoints successfully!') raise #train function with side supervision step=0 writer = SummaryWriter(save_path_cpoints+'summary') best_mae=1 best_epoch=0 clip=0.5 #train function def train(train_loader, model, optimizer, epoch,save_path): global step model.cuda() model.train() loss_all=0 epoch_step=0 try: for i, (images, gts, depths) in enumerate(train_loader, start=1): optimizer.zero_grad() images = images.cuda() gts = gts.cuda() depths=depths.cuda() gts4= nn.functional.interpolate(gts, size=(14,14) , mode='bilinear') gts4=gts4.cuda() gts3= nn.functional.interpolate(gts, size=(28,28) , mode='bilinear') gts3=gts3.cuda() gts2= nn.functional.interpolate(gts, size=(56,56) , mode='bilinear') gts2=gts2.cuda() gts1= nn.functional.interpolate(gts, size=(112,112) , mode='bilinear') gts1=gts1.cuda() ''' #s1,s2 = model(images,depths) s1 = model(images,depths) ''' s1, level1,level2,level3,level4,s1_sig,lev1_sig,lev2_sig,lev3_sig,lev4_sig = model(images,depths) sal_loss= CE(s1,gts) +IOU(s1_sig, gts) loss1= CE(level1,gts1)+IOU(lev1_sig, gts1) loss2= CE(level2,gts2)+IOU(lev2_sig, gts2) loss3= CE(level3,gts3)+IOU(lev3_sig, gts3) loss4= CE(level4,gts4)+IOU(lev4_sig, gts4) loss=sal_loss+loss1+loss2+loss3+loss4 loss.backward() clip_gradient(optimizer, clip) optimizer.step() step+=1 epoch_step+=1 loss_all+=float(loss.data) if i % 100 == 0 or i == total_step or i==1: print('{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], Loss1: {:.4f} '. format(datetime.now(), epoch, epoch, i, total_step, loss.data)) logging.info('#TRAIN#:Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], Loss1: {:.4f} '. format( epoch, epoch, i, total_step, loss.data)) writer.add_scalar('Loss', loss.data, global_step=step) grid_image = make_grid(images[0].clone().cpu().data, 1, normalize=True) writer.add_image('RGB', grid_image, step) grid_image = make_grid(gts[0].clone().cpu().data, 1, normalize=True) writer.add_image('Ground_truth', grid_image, step) res=s1[0].clone() res = res.sigmoid().data.cpu().numpy().squeeze() res = (res - res.min()) / (res.max() - res.min() + 1e-8) writer.add_image('s1', torch.tensor(res), step,dataformats='HW') ''' res=s2[0].clone() res = res.sigmoid().data.cpu().numpy().squeeze() res = (res - res.min()) / (res.max() - res.min() + 1e-8) writer.add_image('s2', torch.tensor(res), step,dataformats='HW') ''' loss_all/=epoch_step logging.info('#TRAIN#:Epoch [{:03d}/{:03d}], Loss_AVG: {:.4f}'.format( epoch, epoch, loss_all)) writer.add_scalar('Loss-epoch', loss_all, global_step=epoch) # if (epoch) % 20 == 0: # torch.save(model.state_dict(), save_path+'IR_Net_epoch_{}.pth'.format(epoch)) except KeyboardInterrupt: #print('Keyboard Interrupt: save model and exit.') #if not os.path.exists(save_path): # os.makedirs(save_path) #torch.save(model.state_dict(), save_path+'IR_Net_epoch_{}.pth'.format(epoch+1)) #print('save checkpoints successfully!') raise #validation test function with edges def test(test_loader,model,epoch,save_path): global best_mae,best_epoch model.eval() with torch.no_grad(): mae_sum=0 for i in range(test_loader.size): image, gt,depth,name,img_for_post = test_loader.load_data() gt = np.asarray(gt, np.float32) gt /= (gt.max() + 1e-8) image = image.cuda() depth = depth.cuda() res,_ = model(image,depth) #_,res = model(image,depth) m = nn.Upsample(size=gt.shape, mode='bilinear', align_corners=False) res=m(res) #res = f.upsample(res, size=gt.shape, mode='bilinear', align_corners=False) res = res.sigmoid().data.cpu().numpy().squeeze() res = (res - res.min()) / (res.max() - res.min() + 1e-8) #imageio.imwrite(save_path+name,res) #misc.imsave(save_path+name, res) mae_sum+=np.sum(np.abs(res-gt))*1.0/(gt.shape[0]*gt.shape[1]) mae=mae_sum/test_loader.size writer.add_scalar('MAE', torch.tensor(mae), global_step=epoch) print('Epoch: {} MAE: {} #### bestMAE: {} bestEpoch: {}'.format(epoch,mae,best_mae,best_epoch)) torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': mae }, '/content/chk_points/current_ckpoint.pt') if epoch==1: best_mae=mae else: if mae