232 lines
9.2 KiB
Python
232 lines
9.2 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
"""
|
|
@Author : Peike Li
|
|
@Contact : peike.li@yahoo.com
|
|
@File : train.py
|
|
@Time : 8/4/19 3:36 PM
|
|
@Desc :
|
|
@License : This source code is licensed under the license found in the
|
|
LICENSE file in the root directory of this source tree.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import timeit
|
|
import argparse
|
|
|
|
import torch
|
|
import torch.optim as optim
|
|
import torchvision.transforms as transforms
|
|
import torch.backends.cudnn as cudnn
|
|
from torch.utils import data
|
|
|
|
import networks
|
|
import utils.schp as schp
|
|
from datasets.datasets import LIPDataSet
|
|
from datasets.target_generation import generate_edge_tensor
|
|
from utils.transforms import BGR2RGB_transform
|
|
from utils.criterion import CriterionAll
|
|
from utils.encoding import DataParallelModel, DataParallelCriterion
|
|
from utils.warmup_scheduler import SGDRScheduler
|
|
|
|
|
|
def get_arguments():
|
|
"""Parse all the arguments provided from the CLI.
|
|
Returns:
|
|
A list of parsed arguments.
|
|
"""
|
|
parser = argparse.ArgumentParser(description="Self Correction for Human Parsing")
|
|
|
|
# Network Structure
|
|
parser.add_argument("--arch", type=str, default='resnet101')
|
|
# Data Preference
|
|
parser.add_argument("--data-dir", type=str, default='./data/LIP')
|
|
parser.add_argument("--batch-size", type=int, default=16)
|
|
parser.add_argument("--input-size", type=str, default='473,473')
|
|
parser.add_argument("--num-classes", type=int, default=20)
|
|
parser.add_argument("--ignore-label", type=int, default=255)
|
|
parser.add_argument("--random-mirror", action="store_true")
|
|
parser.add_argument("--random-scale", action="store_true")
|
|
# Training Strategy
|
|
parser.add_argument("--learning-rate", type=float, default=7e-3)
|
|
parser.add_argument("--momentum", type=float, default=0.9)
|
|
parser.add_argument("--weight-decay", type=float, default=5e-4)
|
|
parser.add_argument("--gpu", type=str, default='0,1,2')
|
|
parser.add_argument("--start-epoch", type=int, default=0)
|
|
parser.add_argument("--epochs", type=int, default=150)
|
|
parser.add_argument("--eval-epochs", type=int, default=10)
|
|
parser.add_argument("--imagenet-pretrain", type=str, default='./pretrain_model/resnet101-imagenet.pth')
|
|
parser.add_argument("--log-dir", type=str, default='./log')
|
|
parser.add_argument("--model-restore", type=str, default='./log/checkpoint.pth.tar')
|
|
parser.add_argument("--schp-start", type=int, default=100, help='schp start epoch')
|
|
parser.add_argument("--cycle-epochs", type=int, default=10, help='schp cyclical epoch')
|
|
parser.add_argument("--schp-restore", type=str, default='./log/schp_checkpoint.pth.tar')
|
|
parser.add_argument("--lambda-s", type=float, default=1, help='segmentation loss weight')
|
|
parser.add_argument("--lambda-e", type=float, default=1, help='edge loss weight')
|
|
parser.add_argument("--lambda-c", type=float, default=0.1, help='segmentation-edge consistency loss weight')
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = get_arguments()
|
|
print(args)
|
|
|
|
start_epoch = 0
|
|
cycle_n = 0
|
|
|
|
if not os.path.exists(args.log_dir):
|
|
os.makedirs(args.log_dir)
|
|
with open(os.path.join(args.log_dir, 'args.json'), 'w') as opt_file:
|
|
json.dump(vars(args), opt_file)
|
|
|
|
gpus = [int(i) for i in args.gpu.split(',')]
|
|
if not args.gpu == 'None':
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
|
|
|
|
input_size = list(map(int, args.input_size.split(',')))
|
|
|
|
cudnn.enabled = True
|
|
cudnn.benchmark = True
|
|
|
|
# Model Initialization
|
|
AugmentCE2P = networks.init_model(args.arch, num_classes=args.num_classes, pretrained=args.imagenet_pretrain)
|
|
model = DataParallelModel(AugmentCE2P)
|
|
model.cuda()
|
|
|
|
IMAGE_MEAN = AugmentCE2P.mean
|
|
IMAGE_STD = AugmentCE2P.std
|
|
INPUT_SPACE = AugmentCE2P.input_space
|
|
print('image mean: {}'.format(IMAGE_MEAN))
|
|
print('image std: {}'.format(IMAGE_STD))
|
|
print('input space:{}'.format(INPUT_SPACE))
|
|
|
|
restore_from = args.model_restore
|
|
if os.path.exists(restore_from):
|
|
print('Resume training from {}'.format(restore_from))
|
|
checkpoint = torch.load(restore_from)
|
|
model.load_state_dict(checkpoint['state_dict'])
|
|
start_epoch = checkpoint['epoch']
|
|
|
|
SCHP_AugmentCE2P = networks.init_model(args.arch, num_classes=args.num_classes, pretrained=args.imagenet_pretrain)
|
|
schp_model = DataParallelModel(SCHP_AugmentCE2P)
|
|
schp_model.cuda()
|
|
|
|
if os.path.exists(args.schp_restore):
|
|
print('Resuming schp checkpoint from {}'.format(args.schp_restore))
|
|
schp_checkpoint = torch.load(args.schp_restore)
|
|
schp_model_state_dict = schp_checkpoint['state_dict']
|
|
cycle_n = schp_checkpoint['cycle_n']
|
|
schp_model.load_state_dict(schp_model_state_dict)
|
|
|
|
# Loss Function
|
|
criterion = CriterionAll(lambda_1=args.lambda_s, lambda_2=args.lambda_e, lambda_3=args.lambda_c,
|
|
num_classes=args.num_classes)
|
|
criterion = DataParallelCriterion(criterion)
|
|
criterion.cuda()
|
|
|
|
# Data Loader
|
|
if INPUT_SPACE == 'BGR':
|
|
print('BGR Transformation')
|
|
transform = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=IMAGE_MEAN,
|
|
std=IMAGE_STD),
|
|
])
|
|
|
|
elif INPUT_SPACE == 'RGB':
|
|
print('RGB Transformation')
|
|
transform = transforms.Compose([
|
|
transforms.ToTensor(),
|
|
BGR2RGB_transform(),
|
|
transforms.Normalize(mean=IMAGE_MEAN,
|
|
std=IMAGE_STD),
|
|
])
|
|
|
|
train_dataset = LIPDataSet(args.data_dir, 'train', crop_size=input_size, transform=transform)
|
|
train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size * len(gpus),
|
|
num_workers=16, shuffle=True, pin_memory=True, drop_last=True)
|
|
print('Total training samples: {}'.format(len(train_dataset)))
|
|
|
|
# Optimizer Initialization
|
|
optimizer = optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum,
|
|
weight_decay=args.weight_decay)
|
|
|
|
lr_scheduler = SGDRScheduler(optimizer, total_epoch=args.epochs,
|
|
eta_min=args.learning_rate / 100, warmup_epoch=10,
|
|
start_cyclical=args.schp_start, cyclical_base_lr=args.learning_rate / 2,
|
|
cyclical_epoch=args.cycle_epochs)
|
|
|
|
total_iters = args.epochs * len(train_loader)
|
|
start = timeit.default_timer()
|
|
for epoch in range(start_epoch, args.epochs):
|
|
lr_scheduler.step(epoch=epoch)
|
|
lr = lr_scheduler.get_lr()[0]
|
|
|
|
model.train()
|
|
for i_iter, batch in enumerate(train_loader):
|
|
i_iter += len(train_loader) * epoch
|
|
|
|
images, labels, _ = batch
|
|
labels = labels.cuda(non_blocking=True)
|
|
|
|
edges = generate_edge_tensor(labels)
|
|
labels = labels.type(torch.cuda.LongTensor)
|
|
edges = edges.type(torch.cuda.LongTensor)
|
|
|
|
preds = model(images)
|
|
|
|
# Online Self Correction Cycle with Label Refinement
|
|
if cycle_n >= 1:
|
|
with torch.no_grad():
|
|
soft_preds = schp_model(images)
|
|
soft_parsing = []
|
|
soft_edge = []
|
|
for soft_pred in soft_preds:
|
|
soft_parsing.append(soft_pred[0][-1])
|
|
soft_edge.append(soft_pred[1][-1])
|
|
soft_preds = torch.cat(soft_parsing, dim=0)
|
|
soft_edges = torch.cat(soft_edge, dim=0)
|
|
else:
|
|
soft_preds = None
|
|
soft_edges = None
|
|
|
|
loss = criterion(preds, [labels, edges, soft_preds, soft_edges], cycle_n)
|
|
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
if i_iter % 100 == 0:
|
|
print('iter = {} of {} completed, lr = {}, loss = {}'.format(i_iter, total_iters, lr,
|
|
loss.data.cpu().numpy()))
|
|
if (epoch + 1) % (args.eval_epochs) == 0:
|
|
schp.save_schp_checkpoint({
|
|
'epoch': epoch + 1,
|
|
'state_dict': model.state_dict(),
|
|
}, False, args.log_dir, filename='checkpoint_{}.pth.tar'.format(epoch + 1))
|
|
|
|
# Self Correction Cycle with Model Aggregation
|
|
if (epoch + 1) >= args.schp_start and (epoch + 1 - args.schp_start) % args.cycle_epochs == 0:
|
|
print('Self-correction cycle number {}'.format(cycle_n))
|
|
schp.moving_average(schp_model, model, 1.0 / (cycle_n + 1))
|
|
cycle_n += 1
|
|
schp.bn_re_estimate(train_loader, schp_model)
|
|
schp.save_schp_checkpoint({
|
|
'state_dict': schp_model.state_dict(),
|
|
'cycle_n': cycle_n,
|
|
}, False, args.log_dir, filename='schp_{}_checkpoint.pth.tar'.format(cycle_n))
|
|
|
|
torch.cuda.empty_cache()
|
|
end = timeit.default_timer()
|
|
print('epoch = {} of {} completed using {} s'.format(epoch, args.epochs,
|
|
(end - start) / (epoch - start_epoch + 1)))
|
|
|
|
end = timeit.default_timer()
|
|
print('Training Finished in {} seconds'.format(end - start))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|