main.py

# -*- coding: utf-8 -*-
import os
import gc
import sys
import copy
import pickle
import logging
import argparse

from utils.trainer_verbose import train_with_ignite, train_without_ignite, get_optimizer
from utils import check_mkdir

import torch

from networks import mobile_hair

logger = logging.getLogger('hair segmentation project')

def str2bool(s):
    return s.lower() in ('t', 'true', '1')


def get_args():
    parser = argparse.ArgumentParser(description='Hair Segmentation')
    parser.add_argument('--networks', default='mobilenet')
    parser.add_argument('--scheduler', default='ReduceLROnPlateau')
    parser.add_argument('--dataset', default='figaro')
    parser.add_argument('--data_dir', default='./data/Figaro1k')
    parser.add_argument('--batch_size', type=int, default=4)
    parser.add_argument('--epochs', default=5, type=int)
    parser.add_argument('--lr', default=0.0001, type=float)
    parser.add_argument('--num_workers', type=int, default=2)
    parser.add_argument('--img_size',type=int, default=256)
    parser.add_argument('--use_pretrained', type=str, default='ImageNet')
    parser.add_argument('--ignite', type=str2bool, default=True)
    parser.add_argument('--visdom', type=str2bool, default=False)
    parser.add_argument('--optimizer', type=str, default='adam')
    parser.add_argument('--momentum', type=float, default=0.9)

    args = parser.parse_args()

    return args


def main():
    args = get_args()
    
    check_mkdir('./logs')
    
    logging_name = './logs/{}_{}_lr_{}.txt'.format(args.networks,
                                                   args.optimizer,
                                                   args.lr)
    
    logger.setLevel(logging.DEBUG)
    formatter = logging.Formatter(
        '[%(asctime)10s][%(levelname)s] %(message)s',
        datefmt='%Y/%m/%d %H:%M:%S'
    )

    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(formatter)

    file_handler = logging.FileHandler(logging_name)
    file_handler.setFormatter(formatter)

    logger.addHandler(stream_handler)
    logger.addHandler(file_handler)
    logger.info('arguments:{}'.format(" ".join(sys.argv)))
    if args.ignite is False:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        
        model = mobile_hair.MobileMattingFCN()
        
        if torch.cuda.is_available():
            if torch.cuda.device_count() > 1:
                print('multi gpu')
                model = torch.nn.DataParallel(model)
        
        model.to(device)
        
        loss = mobile_hair.HairMattingLoss()
        
        optimizer = get_optimizer(args.optimizer, model, args.lr, args.momentum)
        # torch.optim.Adam(filter(lambda p: p.requires_grad,model.parameters()), lr=0.0001, betas=(0.9, 0.999))
        
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min')
        
        train_without_ignite(model, 
                             loss,
                             batch_size=args.batch_size,
                             img_size=args.img_size,
                             epochs=args.epochs,
                             lr=args.lr,
                             num_workers=args.num_workers,
                             optimizer=optimizer,
                             logger=logger,
                             gray_image=True,
                             scheduler=scheduler,
                             viz=args.visdom)
    
    else: train_with_ignite(networks=args.networks,
                      dataset=args.dataset,
                      data_dir=args.data_dir,
                      batch_size=args.batch_size,
                      epochs=args.epochs,
                      lr=args.lr,
                      num_workers=args.num_workers,
                      optimizer=args.optimizer,
                      momentum=args.momentum,
                      img_size=args.img_size,
                      logger=logger)


if __name__ == '__main__':
    main()