train.py

import datetime
import time
import os
import platform

import torch
from torch import nn
from torch import optim
from torch.autograd import Variable
from torch.backends import cudnn
from torch.utils.data import DataLoader
from torchvision import transforms
from tensorboardX import SummaryWriter
from tqdm import tqdm

from settings import getConfig

import joint_transforms
from config import *
from config import backbone_path

from datasets import ImageFolder
from misc import AvgMeter, check_mkdir
from PFNet import PFNet
import settings
import loss

cudnn.benchmark = True

torch.manual_seed(2021)

#Martin CrossPlatform
currOS = platform.system()
cudnn.benchmark = True

#Kaney args parameter
opt = getConfig()
    
print(opt)

device_ids = [opt.device]
ckpt_path = opt.ckpt_path
exp_name = opt.exp_name


args = {
    'epoch_num': opt.epochs,
    'train_batch_size': opt.batch_size,
    'last_epoch': opt.last_epoch,
    'lr': opt.lr,
    'lr_decay': opt.lr_decay,
    'weight_decay': opt.weight_decay,
    'momentum': opt.momentum,
    'snapshot': opt.snap_shot,
    'scale': opt.img_size,
    'save_point': opt.save_point,
    'poly_train': opt.poly_train,
    'optimizer': opt.optimizer,
    
}

print(torch.__version__)

# Path.
check_mkdir(ckpt_path)
check_mkdir(os.path.join(ckpt_path, exp_name))
vis_path = os.path.join(ckpt_path, exp_name, 'log')
check_mkdir(vis_path)
log_path = os.path.join(ckpt_path, exp_name, str(datetime.datetime.now()) + '.txt')
writer = SummaryWriter(log_dir=vis_path, comment=exp_name)

# Transform Data.
joint_transform = joint_transforms.Compose([
    joint_transforms.RandomHorizontallyFlip(),
    joint_transforms.Resize((args['scale'], args['scale']))
])
img_transform = transforms.Compose([
    transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

target_transform = transforms.ToTensor()

# Prepare Data Set.
train_set = ImageFolder(train_path, joint_transform, img_transform, target_transform)


print("Train set: {}".format(train_set.__len__()))
train_loader = DataLoader(train_set, batch_size=args['train_batch_size'], num_workers=opt.num_workers, shuffle=True)

total_epoch = args['epoch_num'] * len(train_loader)



# loss function
#Martin
if currOS == 'Darwin':
    structure_loss = loss.structure_loss().to('mps')
    bce_loss = nn.BCEWithLogitsLoss().to('mps')
    iou_loss = loss.IOU().to('mps')
else:
    structure_loss = loss.structure_loss().cuda(device_ids[0])
    bce_loss = nn.BCEWithLogitsLoss().cuda(device_ids[0])
    iou_loss = loss.IOU().cuda(device_ids[0])
    
    
def bce_iou_loss(pred, target):
    bce_out = bce_loss(pred, target)
    iou_out = iou_loss(pred, target)

    loss = bce_out + iou_out

    return loss

def main():
    print(opt)
    print(args)
    print(exp_name)
    
    
    #Martin - depending on what operating system we are using, choose how to connect to GPU
    if currOS == 'Darwin':
        net = PFNet(backbone_path).to('mps').train()
    else:
        net = PFNet(backbone_path).cuda(device_ids[0]).train() # Kaney modify so it works with PC
        
        
    if args['optimizer'] == 'Adam':
        print("Adam")
        optimizer = optim.Adam([
            {'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
             'lr': 2 * args['lr']},
            {'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
             'lr': 1 * args['lr'], 'weight_decay': args['weight_decay']}
        ])
    else:
        print("SGD")
        optimizer = optim.SGD([
            {'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
             'lr': 2 * args['lr']},
            {'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
             'lr': 1 * args['lr'], 'weight_decay': args['weight_decay']}
        ], momentum=args['momentum'])

    if len(args['snapshot']) > 0:
        print('Training Resumes From \'%s\'' % args['snapshot'])
        net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')))
        total_epoch = (args['epoch_num'] - int(args['snapshot'])) * len(train_loader)
        print(total_epoch)

        
        
    net = nn.DataParallel(net, device_ids=device_ids)
    
    
    print("Using {} GPU(s) to Train.".format(len(device_ids)))

    open(log_path, 'w').write(str(args) + '\n\n')
    train(net, optimizer)
    writer.close()

def train(net, optimizer):
    curr_iter = 1
    start_time = time.time()

    for epoch in range(args['last_epoch'] + 1, args['last_epoch'] + 1 + args['epoch_num']):
        loss_record, loss_1_record, loss_2_record, loss_3_record, loss_4_record = AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()

        train_iterator = tqdm(train_loader, total=len(train_loader))
        for data in train_iterator:
            if args['poly_train']:
                base_lr = args['lr'] * (1 - float(curr_iter) / float(total_epoch)) ** args['lr_decay']
                optimizer.param_groups[0]['lr'] = 2 * base_lr
                optimizer.param_groups[1]['lr'] = 1 * base_lr

            inputs, labels = data
            batch_size = inputs.size(0)
            #Martin
            if currOS == 'Darwin':
                inputs = Variable(inputs).to('mps')
                labels = Variable(labels).to('mps')
            else:
                inputs = Variable(inputs).cuda(device_ids[0])
                labels = Variable(labels).cuda(device_ids[0])

            optimizer.zero_grad()

            predict_1, predict_2, predict_3, predict_4 = net(inputs)

            loss_1 = bce_iou_loss(predict_1, labels)
            loss_2 = structure_loss(predict_2, labels)
            loss_3 = structure_loss(predict_3, labels)
            loss_4 = structure_loss(predict_4, labels)

            loss = 1 * loss_1 + 1 * loss_2 + 2 * loss_3 + 4 * loss_4

            loss.backward()

            optimizer.step()

            loss_record.update(loss.data, batch_size)
            loss_1_record.update(loss_1.data, batch_size)
            loss_2_record.update(loss_2.data, batch_size)
            loss_3_record.update(loss_3.data, batch_size)
            loss_4_record.update(loss_4.data, batch_size)

            if curr_iter % 10 == 0:
                writer.add_scalar('loss', loss, curr_iter)
                writer.add_scalar('loss_1', loss_1, curr_iter)
                writer.add_scalar('loss_2', loss_2, curr_iter)
                writer.add_scalar('loss_3', loss_3, curr_iter)
                writer.add_scalar('loss_4', loss_4, curr_iter)

            log = '[%3d], [%6d], [%.6f], [%.5f], [%.5f], [%.5f], [%.5f], [%.5f]' % \
                  (epoch, curr_iter, base_lr, loss_record.avg, loss_1_record.avg, loss_2_record.avg,
                   loss_3_record.avg, loss_4_record.avg)
            train_iterator.set_description(log)
            open(log_path, 'a').write(log + '\n')

            curr_iter += 1

        if epoch in args['save_point']:
            net.cpu()
            torch.save(net.module.state_dict(), os.path.join(ckpt_path, exp_name, '%d.pth' % epoch))
               #Martin
            if currOS == 'Darwin':
                net.to('mps')
            else:
                net.cuda(device_ids[0])
                
        if epoch >= args['epoch_num']:
            net.cpu()
            torch.save(net.module.state_dict(), os.path.join(ckpt_path, exp_name, '%d.pth' % epoch))
            print("Total Training Time: {}".format(str(datetime.timedelta(seconds=int(time.time() - start_time)))))
            print(exp_name)
            print("Optimization Have Done!")
            return

if __name__ == '__main__':
    main()