train.py

import torch
import torchvision
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torch.optim as optim
import torchvision.transforms as standard_transforms

import numpy as np
import glob

from data_loader import Rescale
from data_loader import RescaleT
from data_loader import RandomCrop
from data_loader import ToTensor
from data_loader import ToTensorLab
from data_loader import SalObjDataset

from models import U2NET
from models import U2NETP
from models import JITNET

# ------- 1. define loss function --------

bce_loss = nn.BCELoss(size_average=True)

def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v):

  loss0 = bce_loss(d0,labels_v)
  loss1 = bce_loss(d1,labels_v)
  loss2 = bce_loss(d2,labels_v)
  loss3 = bce_loss(d3,labels_v)
  loss4 = bce_loss(d4,labels_v)
  loss5 = bce_loss(d5,labels_v)
  loss6 = bce_loss(d6,labels_v)

  loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6
  print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n"%(loss0.data[0],loss1.data[0],loss2.data[0],loss3.data[0],loss4.data[0],loss5.data[0],loss6.data[0]))

  return loss0, loss


# ------- 2. set the directory of training dataset --------

model_name = 'jitnet' #'u2netp'

data_dir = './data/'
tra_image_dir = 'supervisely/'
# tra_label_dir = 'portrait/img_crop/'

image_ext = '.png'
label_ext = '.png'

model_dir = './saved_models/' + model_name +'/'

epoch_num = 100000
batch_size_train = 12
batch_size_val = 1
train_num = 0
val_num = 0

print(data_dir + tra_image_dir + '**/img/*' + image_ext)

tra_img_name_list = glob.glob(data_dir + tra_image_dir + '**/img/*' + image_ext)

tra_lbl_name_list = [ img_path.replace('img', 'masks_machine') for img_path in tra_img_name_list]
# for img_path in tra_img_name_list:
  # img_name = img_path.split("/")[-1]

  # aaa = img_name.split(".")
  # bbb = aaa[0:-1]
  # imidx = bbb[0]
  # for i in range(1,len(bbb)):
  #   imidx = imidx + "." + bbb[i]

  # tra_lbl_name_list.append(data_dir + tra_label_dir + imidx + label_ext)

print("---")
print("train images: ", len(tra_img_name_list))
print("train labels: ", len(tra_lbl_name_list))
print("---")

train_num = len(tra_img_name_list)

salobj_dataset = SalObjDataset(
    img_name_list=tra_img_name_list,
    lbl_name_list=tra_lbl_name_list,
    transform=transforms.Compose([
        RescaleT(320),
        RandomCrop(288),
        ToTensorLab(flag=0)]))
salobj_dataloader = DataLoader(salobj_dataset, batch_size=batch_size_train, shuffle=True, num_workers=8)

# ------- 3. define model --------
# define the net
if model_name=='u2net':
    net = U2NET(3, 1)
elif model_name=='u2netp':
    net = U2NETP(3,1)
elif model_name=='jitnet':
    net = JITNET(3, 1)


if torch.cuda.is_available():
    net.cuda()

# ------- 4. define optimizer --------
print("---define optimizer...")
optimizer = optim.Adam(net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)

# ------- 5. training process --------
print("---start training...")
ite_num = 0
running_loss = 0.0
running_tar_loss = 0.0
ite_num4val = 0
save_frq = 1000 # save the model every 1000 iterations

for epoch in range(0, epoch_num):
    net.train()

    for i, data in enumerate(salobj_dataloader):
        ite_num = ite_num + 1
        ite_num4val = ite_num4val + 1

        inputs, labels = data['image'], data['label']

        inputs = inputs.type(torch.FloatTensor)
        labels = labels.type(torch.FloatTensor)

        # wrap them in Variable
        if torch.cuda.is_available():
            inputs_v, labels_v = Variable(inputs.cuda(), requires_grad=False), Variable(labels.cuda(),
                                                                                        requires_grad=False)
        else:
            inputs_v, labels_v = Variable(inputs, requires_grad=False), Variable(labels, requires_grad=False)

        # y zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        d0, d1, d2, d3, d4, d5, d6 = net(inputs_v)
        if model_name == 'jitnet':
          loss = bce_loss(d0, labels_v)
          loss2 = loss
        else:
          loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v)

        loss.backward()
        optimizer.step()

        # # print statistics
        running_loss += loss.item()
        running_tar_loss += loss2.item()

        # del temporary outputs and loss
        del d0, d1, d2, d3, d4, d5, d6, loss, loss2

        print("[epoch: %3d/%3d, batch: %5d/%5d, ite: %d] train loss: %3f, tar: %3f " % (
        epoch + 1, epoch_num, (i + 1) * batch_size_train, train_num, ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val))

        if ite_num % save_frq == 0:

            torch.save(net.state_dict(), model_dir + model_name+"_bce_itr_%d_train_%3f_tar_%3f.pth" % (ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val))
            running_loss = 0.0
            running_tar_loss = 0.0
            net.train()  # resume train
            ite_num4val = 0

if __name__ == "__main__":
    main()