def main():
parser = argparse.ArgumentParser()
parser.add_argument('-e', '--exp_name', default='fcn32s_color')
parser.add_argument('-g', '--gpu', type=int, required=True)
parser.add_argument('-c',
'--config',
type=int,
default=1,
choices=configurations.keys())
parser.add_argument('-b',
'--binning',
default='soft',
choices=('soft', 'one-hot', 'uniform'))
parser.add_argument('-k', '--numbins', type=int, default=128)
parser.add_argument(
'-d',
'--dataset_path',
default='/vis/home/arunirc/data1/datasets/ImageNet/images/')
parser.add_argument('-m', '--model_path', default=None)
parser.add_argument('--resume', help='Checkpoint path')
args = parser.parse_args()
# -----------------------------------------------------------------------------
# 0. setup
# -----------------------------------------------------------------------------
gpu = args.gpu
cfg = configurations[args.config]
cfg.update({'bin_type': args.binning, 'numbins': args.numbins})
resume = args.resume
if resume:
out, _ = osp.split(resume)
else:
out = get_log_dir(args.exp_name, args.config, cfg, verbose=False)
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
torch.cuda.manual_seed(1337)
torch.backends.cudnn.enabled = True
# torch.backends.cudnn.benchmark = True
# -----------------------------------------------------------------------------
# 1. dataset
# -----------------------------------------------------------------------------
# Custom dataset class defined in `data_loader.ColorizeImageNet`
root = args.dataset_path
kwargs = {'num_workers': 4, 'pin_memory': True} if cuda else {}
if 'img_lowpass' in cfg.keys():
img_lowpass = cfg['img_lowpass']
else:
img_lowpass = None
if 'train_set' in cfg.keys():
train_set = cfg['train_set']
else:
train_set = 'full'
if 'val_set' in cfg.keys():
val_set = cfg['val_set']
else:
val_set = 'small'
if 'gmm_path' in cfg.keys():
gmm_path = cfg['gmm_path']
else:
gmm_path = None
if 'mean_l_path' in cfg.keys():
mean_l_path = cfg['mean_l_path']
else:
mean_l_path = None
if 'im_size' in cfg.keys():
im_size = cfg['im_size']
else:
im_size = (256, 256)
if 'batch_size' in cfg.keys():
batch_size = cfg['batch_size']
else:
batch_size = 1
if 'uniform_sigma' in cfg.keys():
uniform_sigma = cfg['uniform_sigma']
else:
uniform_sigma = 'default'
if 'binning' in cfg.keys():
args.binning = cfg['binning']
# DEBUG: set='tiny'
train_loader = torch.utils.data.DataLoader(
data_loader.ColorizeImageNet(root,
split='train',
bins=args.binning,
log_dir=out,
num_hc_bins=args.numbins,
set=train_set,
img_lowpass=img_lowpass,
im_size=im_size,
gmm_path=gmm_path,
mean_l_path=mean_l_path,
uniform_sigma=uniform_sigma),
batch_size=batch_size,
shuffle=True,
**kwargs) # DEBUG: set shuffle False
# DEBUG: set='tiny'
val_loader = torch.utils.data.DataLoader(data_loader.ColorizeImageNet(
root,
split='val',
bins=args.binning,
log_dir=out,
num_hc_bins=args.numbins,
set=val_set,
img_lowpass=img_lowpass,
im_size=im_size,
gmm_path=gmm_path,
mean_l_path=mean_l_path,
uniform_sigma=uniform_sigma),
batch_size=1,
shuffle=False,
**kwargs)
# -----------------------------------------------------------------------------
# 2. model
# -----------------------------------------------------------------------------
model = models.FCN32sColor(n_class=args.numbins, bin_type=args.binning)
if args.model_path:
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['model_state_dict'])
start_epoch = 0
start_iteration = 0
if resume:
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['model_state_dict'])
start_epoch = checkpoint['epoch']
start_iteration = checkpoint['iteration']
else:
pass
if cuda:
model = model.cuda()
# -----------------------------------------------------------------------------
# 3. optimizer
# -----------------------------------------------------------------------------
params = filter(lambda p: p.requires_grad, model.parameters())
if 'optim' in cfg.keys():
if cfg['optim'].lower() == 'sgd':
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
elif cfg['optim'].lower() == 'adam':
optim = torch.optim.Adam(params,
lr=cfg['lr'],
weight_decay=cfg['weight_decay'])
else:
raise NotImplementedError('Optimizers: SGD or Adam')
else:
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
if resume:
optim.load_state_dict(checkpoint['optim_state_dict'])
# -----------------------------------------------------------------------------
# Sanity-check: forward pass with a single sample
# -----------------------------------------------------------------------------
DEBUG = False
if DEBUG:
dataiter = iter(val_loader)
img, label = dataiter.next()
model.eval()
print 'Labels: ' + str(label.size()) # batchSize x num_class
print 'Input: ' + str(img.size()) # batchSize x 1 x (im_size)
if val_loader.dataset.bins == 'one-hot':
from torch.autograd import Variable
inputs = Variable(img)
if cuda:
inputs = inputs.cuda()
outputs = model(inputs)
assert len(outputs)==2, \
'Network should predict a 2-tuple: hue-map and chroma-map.'
hue_map = outputs[0].data
chroma_map = outputs[1].data
assert hue_map.size() == chroma_map.size(), \
'Outputs should have same dimensions.'
sz_h = hue_map.size()
sz_im = img.size()
assert sz_im[2]==sz_h[2] and sz_im[3]==sz_h[3], \
'Spatial dims should match for input and output.'
elif val_loader.dataset.bins == 'soft':
from torch.autograd import Variable
inputs = Variable(img)
if cuda:
inputs = inputs.cuda()
outputs = model(inputs)
# TODO: assertions
# del inputs, outputs
import pdb
pdb.set_trace() # breakpoint 0632fd52 //
model.train()
else:
pass
# -----------------------------------------------------------------------------
# Training
# -----------------------------------------------------------------------------
trainer = train.Trainer(
cuda=cuda,
model=model,
optimizer=optim,
train_loader=train_loader,
val_loader=val_loader,
out=out,
max_iter=cfg['max_iteration'],
interval_validate=cfg.get('interval_validate', len(train_loader)),
)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-e', '--exp_name', default='resnet50_vggface')
parser.add_argument('-c', '--config', type=int, default=1,
choices=configurations.keys())
parser.add_argument('-d', '--dataset_path',
default='/srv/data1/arunirc/datasets/vggface2')
parser.add_argument('-m', '--model_path', default=None,
help='Initialize from pre-trained model')
parser.add_argument('--resume', help='Checkpoint path')
parser.add_argument('--bottleneck', action='store_true', default=False,
help='Add a 512-dim bottleneck layer with L2 normalization')
args = parser.parse_args()
# gpu = args.gpu
cfg = configurations[args.config]
out = get_log_dir(args.exp_name, args.config, cfg, verbose=False)
resume = args.resume
# os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
torch.cuda.manual_seed(1337)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True # enable if all images are same size
# -----------------------------------------------------------------------------
# 1. Dataset
# -----------------------------------------------------------------------------
# Images should be arranged like this:
# data_root/
# class_1/....jpg..
# class_2/....jpg..
# ......./....jpg..
data_root = args.dataset_path
kwargs = {'num_workers': 4, 'pin_memory': True} if cuda else {}
RGB_MEAN = [ 0.485, 0.456, 0.406 ]
RGB_STD = [ 0.229, 0.224, 0.225 ]
# Data transforms
# http://pytorch.org/docs/master/torchvision/transforms.html
train_transform = transforms.Compose([
transforms.Scale(256), # smaller side resized
transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean = RGB_MEAN,
std = RGB_STD),
])
val_transform = transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean = RGB_MEAN,
std = RGB_STD),
])
# Data loaders - using PyTorch built-in objects
# loader = DataLoaderClass(DatasetClass)
# * `DataLoaderClass` is PyTorch provided torch.utils.data.DataLoader
# * `DatasetClass` loads samples from a dataset; can be a standard class
# provided by PyTorch (datasets.ImageFolder) or a custom-made class.
# - More info: http://pytorch.org/docs/master/torchvision/datasets.html#imagefolder
traindir = osp.join(data_root, 'train')
dataset_train = datasets.ImageFolder(traindir, train_transform)
# For unbalanced dataset we create a weighted sampler
# * Balanced class sampling: https://discuss.pytorch.org/t/balanced-sampling-between-classes-with-torchvision-dataloader/2703/3
weights = utils.make_weights_for_balanced_classes(
dataset_train.imgs, len(dataset_train.classes))
weights = torch.DoubleTensor(weights)
sampler = torch.utils.data.sampler.WeightedRandomSampler(weights, len(weights))
train_loader = torch.utils.data.DataLoader(
dataset_train, batch_size=cfg['batch_size'],
sampler = sampler, **kwargs)
valdir = osp.join(data_root, 'val-crop')
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(valdir, val_transform),
batch_size=cfg['batch_size'], shuffle=False, **kwargs)
# print 'dataset classes:' + str(train_loader.dataset.classes)
num_class = len(train_loader.dataset.classes)
print 'Number of classes: %d' % num_class
# -----------------------------------------------------------------------------
# 2. Model
# -----------------------------------------------------------------------------
model = torchvision.models.resnet50(pretrained=False)
if type(model.fc) == torch.nn.modules.linear.Linear:
# Check if final fc layer sizes match num_class
if not model.fc.weight.size()[0] == num_class:
# Replace last layer
print model.fc
model.fc = torch.nn.Linear(2048, num_class)
print model.fc
else:
pass
else:
pass
if args.model_path:
# If existing model is to be loaded from a file
checkpoint = torch.load(args.model_path)
if checkpoint['arch'] == 'DataParallel':
# if we trained and saved our model using DataParallel
model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3, 4, 5, 6, 7])
model.load_state_dict(checkpoint['model_state_dict'])
model = model.module # get network module from inside its DataParallel wrapper
else:
model.load_state_dict(checkpoint['model_state_dict'])
# Optionally add a "bottleneck + L2-norm" layer after GAP-layer
# TODO -- loading a bottleneck model might be a problem .... do some unit-tests
if args.bottleneck:
layers = []
layers.append(torch.nn.Linear(2048, 512))
layers.append(nn.BatchNorm2d(512))
layers.append(torch.nn.ReLU(inplace=True))
layers.append(models.NormFeat()) # L2-normalization layer
layers.append(torch.nn.Linear(512, num_class))
model.fc = torch.nn.Sequential(*layers)
# TODO - config options for DataParallel and device_ids
model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3, 4, 5, 6, 7])
if cuda:
model.cuda()
start_epoch = 0
start_iteration = 0
# Loss - cross entropy between predicted scores (unnormalized) and class labels (integers)
criterion = nn.CrossEntropyLoss()
if cuda:
criterion = criterion.cuda()
if resume:
# Resume training from last saved checkpoint
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['model_state_dict'])
start_epoch = checkpoint['epoch']
start_iteration = checkpoint['iteration']
else:
pass
# -----------------------------------------------------------------------------
# 3. Optimizer
# -----------------------------------------------------------------------------
params = filter(lambda p: p.requires_grad, model.parameters())
# Parameters with p.requires_grad=False are not updated during training.
# This can be specified when defining the nn.Modules during model creation
if 'optim' in cfg.keys():
if cfg['optim'].lower()=='sgd':
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
elif cfg['optim'].lower()=='adam':
optim = torch.optim.Adam(params,
lr=cfg['lr'], weight_decay=cfg['weight_decay'])
else:
raise NotImplementedError('Optimizers: SGD or Adam')
else:
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
if resume:
optim.load_state_dict(checkpoint['optim_state_dict'])
# -----------------------------------------------------------------------------
# [optional] Sanity-check: forward pass with a single batch
# -----------------------------------------------------------------------------
DEBUG = False
if DEBUG:
# model = model.cpu()
dataiter = iter(val_loader)
img, label = dataiter.next()
print 'Labels: ' + str(label.size()) # batchSize x num_class
print 'Input: ' + str(img.size()) # batchSize x 3 x 224 x 224
im = img.squeeze().numpy()
im = im[0,:,:,:] # get first image in the batch
im = im.transpose((1,2,0)) # permute to 224x224x3
im = im * [ 0.229, 0.224, 0.225 ] # unnormalize
im = im + [ 0.485, 0.456, 0.406 ]
im[im<0] = 0
f = plt.figure()
plt.imshow(im)
plt.savefig('sanity-check-im.jpg') # save transformed image in current folder
inputs = Variable(img)
if cuda:
inputs = inputs.cuda()
model.eval()
outputs = model(inputs)
print 'Network output: ' + str(outputs.size())
model.train()
else:
pass
# -----------------------------------------------------------------------------
# 4. Training
# -----------------------------------------------------------------------------
trainer = train.Trainer(
cuda=cuda,
model=model,
criterion=criterion,
optimizer=optim,
init_lr=cfg['lr'],
lr_decay_epoch = cfg['lr_decay_epoch'],
train_loader=train_loader,
val_loader=val_loader,
out=out,
max_iter=cfg['max_iteration'],
interval_validate=cfg.get('interval_validate', len(train_loader)),
)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-g', '--gpu', type=int, required=True)
parser.add_argument('-c', '--config', type=int, default=15,
choices=configurations.keys())
parser.add_argument('-b', '--binning', default='uniform',
choices=('soft','one-hot', 'uniform'))
parser.add_argument('-k', '--numbins', type=int, default=128)
parser.add_argument('-d', '--dataset_path',
default='/vis/home/arunirc/data1/datasets/ImageNet/images/')
parser.add_argument('-m', '--model_path', default=None)
parser.add_argument('--data_par', action='store_true', default=False,
help='Use DataParallel for multi-gpu training')
parser.add_argument('--resume', help='Checkpoint path')
args = parser.parse_args()
gpu = args.gpu
# The config must specify path to pre-trained model as value for the
# key 'fcn16s_pretrained_model'
cfg = configurations[args.config]
cfg.update({'bin_type':args.binning,'numbins':args.numbins})
out = get_log_dir('fcn8s_color', args.config, cfg, verbose=False)
resume = args.resume
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
torch.cuda.manual_seed(1337)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
# -----------------------------------------------------------------------------
# 1. dataset
# -----------------------------------------------------------------------------
# root = osp.expanduser('~/data/datasets')
root = args.dataset_path
kwargs = {'num_workers': 4, 'pin_memory': True} if cuda else {}
if 'img_lowpass' in cfg.keys():
img_lowpass = cfg['img_lowpass']
else:
img_lowpass = None
if 'train_set' in cfg.keys():
train_set = cfg['train_set']
else:
train_set = 'full'
if 'val_set' in cfg.keys():
val_set = cfg['val_set']
else:
val_set = 'small'
if 'gmm_path' in cfg.keys():
gmm_path = cfg['gmm_path']
else:
gmm_path = None
if 'mean_l_path' in cfg.keys():
mean_l_path = cfg['mean_l_path']
else:
mean_l_path = None
if 'im_size' in cfg.keys():
im_size = cfg['im_size']
else:
im_size = (256, 256)
if 'batch_size' in cfg.keys():
batch_size = cfg['batch_size']
else:
batch_size = 1
if 'uniform_sigma' in cfg.keys():
uniform_sigma = cfg['uniform_sigma']
else:
uniform_sigma = 'default'
if 'binning' in cfg.keys():
args.binning = cfg['binning']
# DEBUG: set='tiny'
train_loader = torch.utils.data.DataLoader(
data_loader.ColorizeImageNet(root, split='train',
bins=args.binning, log_dir=out, num_hc_bins=args.numbins,
set=train_set, img_lowpass=img_lowpass, im_size=im_size,
gmm_path=gmm_path, mean_l_path=mean_l_path, uniform_sigma=uniform_sigma ),
batch_size=24, shuffle=True, **kwargs) # DEBUG: set shuffle False
# DEBUG: set='tiny'
val_loader = torch.utils.data.DataLoader(
data_loader.ColorizeImageNet(root, split='val',
bins=args.binning, log_dir=out, num_hc_bins=args.numbins,
set=val_set, img_lowpass=img_lowpass, im_size=im_size,
gmm_path=gmm_path, mean_l_path=mean_l_path, uniform_sigma=uniform_sigma ),
batch_size=1, shuffle=False, **kwargs)
# -----------------------------------------------------------------------------
# 2. model
# -----------------------------------------------------------------------------
model = models.FCN8sColor(n_class=args.numbins, bin_type=args.binning)
if args.model_path:
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['model_state_dict'])
else:
if resume:
# HACK: takes very long ... better to start a new expt with init from `args.model_path`
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['model_state_dict'])
start_epoch = checkpoint['epoch']
start_iteration = checkpoint['iteration']
else:
fcn16s = models.FCN16sColor(n_class=args.numbins, bin_type=args.binning)
fcn16s.load_state_dict(torch.load(cfg['fcn16s_pretrained_model'])['model_state_dict'])
model.copy_params_from_fcn16s(fcn16s)
start_epoch = 0
start_iteration = 0
if cuda:
model = model.cuda()
if args.data_par:
raise NotImplementedError
# model = torch.nn.DataParallel(model, device_ids=[1, 2, 3, 4, 5, 6])
# -----------------------------------------------------------------------------
# 3. optimizer
# -----------------------------------------------------------------------------
params = filter(lambda p: p.requires_grad, model.parameters())
if 'optim' in cfg.keys():
if cfg['optim'].lower()=='sgd':
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
elif cfg['optim'].lower()=='adam':
optim = torch.optim.Adam(params,
lr=cfg['lr'], weight_decay=cfg['weight_decay'])
else:
raise NotImplementedError('Optimizers: SGD or Adam')
else:
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
if resume:
optim.load_state_dict(checkpoint['optim_state_dict'])
# -----------------------------------------------------------------------------
# Sanity-check: forward pass with a single sample
# -----------------------------------------------------------------------------
# dataiter = iter(val_loader)
# img, label = dataiter.next()
# model.eval()
# if val_loader.dataset.bins == 'one-hot':
# from torch.autograd import Variable
# inputs = Variable(img)
# if cuda:
# inputs = inputs.cuda()
# outputs = model(inputs)
# assert len(outputs)==2, \
# 'Network should predict a 2-tuple: hue-map and chroma-map.'
# hue_map = outputs[0].data
# chroma_map = outputs[1].data
# assert hue_map.size() == chroma_map.size(), \
# 'Outputs should have same dimensions.'
# sz_h = hue_map.size()
# sz_im = img.size()
# assert sz_im[2]==sz_h[2] and sz_im[3]==sz_h[3], \
# 'Spatial dims should match for input and output.'
# elif val_loader.dataset.bins == 'soft':
# from torch.autograd import Variable
# inputs = Variable(img)
# if cuda:
# inputs = inputs.cuda()
# outputs = model(inputs)
# # TODO: assertions
# # del inputs, outputs
# model.train()
# -----------------------------------------------------------------------------
# Training
# -----------------------------------------------------------------------------
trainer = train.Trainer(
cuda=cuda,
model=model,
optimizer=optim,
train_loader=train_loader,
val_loader=val_loader,
out=out,
max_iter=cfg['max_iteration'],
interval_validate=cfg.get('interval_validate', len(train_loader)),
)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-e', '--exp_name', default='resnet50_vggface')
parser.add_argument('-c',
'--config',
type=int,
default=1,
choices=configurations.keys())
parser.add_argument('-d',
'--dataset_path',
default='/srv/data1/arunirc/datasets/vggface2')
parser.add_argument('-m',
'--model_path',
default=None,
help='Initialize from pre-trained model')
parser.add_argument('--resume', help='Checkpoint path')
parser.add_argument(
'--bottleneck',
action='store_true',
default=False,
help='Add a 512-dim bottleneck layer with L2 normalization')
args = parser.parse_args()
# gpu = args.gpu
cfg = configurations[args.config]
out = get_log_dir(args.exp_name, args.config, cfg, verbose=False)
resume = args.resume
# os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
torch.cuda.manual_seed(1337)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True # enable if all images are same size
# -----------------------------------------------------------------------------
# 1. Dataset
# -----------------------------------------------------------------------------
# Images should be arranged like this:
# data_root/
# class_1/....jpg..
# class_2/....jpg..
# ......./....jpg..
data_root = args.dataset_path
kwargs = {'num_workers': 4, 'pin_memory': True} if cuda else {}
RGB_MEAN = [0.485, 0.456, 0.406]
RGB_STD = [0.229, 0.224, 0.225]
# Data transforms
# http://pytorch.org/docs/master/torchvision/transforms.html
train_transform = transforms.Compose([
transforms.Scale(256), # smaller side resized
transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=RGB_MEAN, std=RGB_STD),
])
val_transform = transforms.Compose([
transforms.Scale(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=RGB_MEAN, std=RGB_STD),
])
# Data loaders - using PyTorch built-in objects
# loader = DataLoaderClass(DatasetClass)
# * `DataLoaderClass` is PyTorch provided torch.utils.data.DataLoader
# * `DatasetClass` loads samples from a dataset; can be a standard class
# provided by PyTorch (datasets.ImageFolder) or a custom-made class.
# - More info: http://pytorch.org/docs/master/torchvision/datasets.html#imagefolder
traindir = osp.join(data_root, 'train')
dataset_train = datasets.ImageFolder(traindir, train_transform)
# For unbalanced dataset we create a weighted sampler
# * Balanced class sampling: https://discuss.pytorch.org/t/balanced-sampling-between-classes-with-torchvision-dataloader/2703/3
weights = utils.make_weights_for_balanced_classes(
dataset_train.imgs, len(dataset_train.classes))
weights = torch.DoubleTensor(weights)
sampler = torch.utils.data.sampler.WeightedRandomSampler(
weights, len(weights))
train_loader = torch.utils.data.DataLoader(dataset_train,
batch_size=cfg['batch_size'],
sampler=sampler,
**kwargs)
valdir = osp.join(data_root, 'val-crop')
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir, val_transform),
batch_size=cfg['batch_size'],
shuffle=False,
**kwargs)
# print 'dataset classes:' + str(train_loader.dataset.classes)
num_class = len(train_loader.dataset.classes)
print 'Number of classes: %d' % num_class
# -----------------------------------------------------------------------------
# 2. Model
# -----------------------------------------------------------------------------
model = torchvision.models.resnet50(pretrained=False)
if type(model.fc) == torch.nn.modules.linear.Linear:
# Check if final fc layer sizes match num_class
if not model.fc.weight.size()[0] == num_class:
# Replace last layer
print model.fc
model.fc = torch.nn.Linear(2048, num_class)
print model.fc
else:
pass
else:
pass
if args.model_path:
# If existing model is to be loaded from a file
checkpoint = torch.load(args.model_path)
if checkpoint['arch'] == 'DataParallel':
# if we trained and saved our model using DataParallel
model = torch.nn.DataParallel(model,
device_ids=[0, 1, 2, 3, 4, 5, 6, 7])
model.load_state_dict(checkpoint['model_state_dict'])
model = model.module # get network module from inside its DataParallel wrapper
else:
model.load_state_dict(checkpoint['model_state_dict'])
# Optionally add a "bottleneck + L2-norm" layer after GAP-layer
# TODO -- loading a bottleneck model might be a problem .... do some unit-tests
if args.bottleneck:
layers = []
layers.append(torch.nn.Linear(2048, 512))
layers.append(nn.BatchNorm2d(512))
layers.append(torch.nn.ReLU(inplace=True))
layers.append(models.NormFeat()) # L2-normalization layer
layers.append(torch.nn.Linear(512, num_class))
model.fc = torch.nn.Sequential(*layers)
# TODO - config options for DataParallel and device_ids
model = torch.nn.DataParallel(model, device_ids=[0, 1, 2, 3, 4, 5, 6, 7])
if cuda:
model.cuda()
start_epoch = 0
start_iteration = 0
# Loss - cross entropy between predicted scores (unnormalized) and class labels (integers)
criterion = nn.CrossEntropyLoss()
if cuda:
criterion = criterion.cuda()
if resume:
# Resume training from last saved checkpoint
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['model_state_dict'])
start_epoch = checkpoint['epoch']
start_iteration = checkpoint['iteration']
else:
pass
# -----------------------------------------------------------------------------
# 3. Optimizer
# -----------------------------------------------------------------------------
params = filter(lambda p: p.requires_grad, model.parameters())
# Parameters with p.requires_grad=False are not updated during training.
# This can be specified when defining the nn.Modules during model creation
if 'optim' in cfg.keys():
if cfg['optim'].lower() == 'sgd':
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
elif cfg['optim'].lower() == 'adam':
optim = torch.optim.Adam(params,
lr=cfg['lr'],
weight_decay=cfg['weight_decay'])
else:
raise NotImplementedError('Optimizers: SGD or Adam')
else:
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
if resume:
optim.load_state_dict(checkpoint['optim_state_dict'])
# -----------------------------------------------------------------------------
# [optional] Sanity-check: forward pass with a single batch
# -----------------------------------------------------------------------------
DEBUG = False
if DEBUG:
# model = model.cpu()
dataiter = iter(val_loader)
img, label = dataiter.next()
print 'Labels: ' + str(label.size()) # batchSize x num_class
print 'Input: ' + str(img.size()) # batchSize x 3 x 224 x 224
im = img.squeeze().numpy()
im = im[0, :, :, :] # get first image in the batch
im = im.transpose((1, 2, 0)) # permute to 224x224x3
im = im * [0.229, 0.224, 0.225] # unnormalize
im = im + [0.485, 0.456, 0.406]
im[im < 0] = 0
f = plt.figure()
plt.imshow(im)
plt.savefig(
'sanity-check-im.jpg') # save transformed image in current folder
inputs = Variable(img)
if cuda:
inputs = inputs.cuda()
model.eval()
outputs = model(inputs)
print 'Network output: ' + str(outputs.size())
model.train()
else:
pass
# -----------------------------------------------------------------------------
# 4. Training
# -----------------------------------------------------------------------------
trainer = train.Trainer(
cuda=cuda,
model=model,
criterion=criterion,
optimizer=optim,
init_lr=cfg['lr'],
lr_decay_epoch=cfg['lr_decay_epoch'],
train_loader=train_loader,
val_loader=val_loader,
out=out,
max_iter=cfg['max_iteration'],
interval_validate=cfg.get('interval_validate', len(train_loader)),
)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-e', '--exp_name', default='resnet_face_demo')
parser.add_argument('-g', '--gpu', type=int, required=True)
parser.add_argument('-c',
'--config',
type=int,
default=1,
choices=configurations.keys())
parser.add_argument('-d',
'--dataset_path',
default='./samples/tiny_dataset')
parser.add_argument('-m', '--model_path', default=None)
parser.add_argument('--resume', help='Checkpoint path')
args = parser.parse_args()
gpu = args.gpu
cfg = configurations[args.config]
out = get_log_dir(args.exp_name, args.config, cfg, verbose=False)
resume = args.resume
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
torch.cuda.manual_seed(1337)
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True # enable if all images are same size
# -----------------------------------------------------------------------------
# 1. Dataset
# -----------------------------------------------------------------------------
# Images should be arranged like this:
# data_root/
# class_1/....jpg..
# class_2/....jpg..
# ......./....jpg..
data_root = args.dataset_path
kwargs = {'num_workers': 4, 'pin_memory': True} if cuda else {}
# Data transforms
# http://pytorch.org/docs/master/torchvision/transforms.html
transform = transforms.Compose([
transforms.Scale(256), # smallest size resizes to 256
transforms.CenterCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
# Data loader
# http://pytorch.org/docs/master/torchvision/datasets.html#imagefolder
traindir = args.dataset_path
train_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
traindir, transform),
batch_size=3,
shuffle=True,
**kwargs)
# for demo purpose, set to be same as train
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
traindir, transform),
batch_size=2,
shuffle=False,
**kwargs)
print 'dataset classes:' + str(train_loader.dataset.classes)
num_class = len(train_loader.dataset.classes)
# -----------------------------------------------------------------------------
# 2. Model
# -----------------------------------------------------------------------------
# PyTorch ResNet model definition:
# https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
# ResNet docs:
# http://pytorch.org/docs/master/torchvision/models.html#id3
model = torchvision.models.resnet50(
pretrained=True) # Using pre-trained for demo purpose
# by default, resnet has 1000 output categories
print model.fc # Check: Linear (2048 -> 1000)
model.fc = torch.nn.Linear(
2048, num_class) # change to current dataset's classes
print model.fc
if args.model_path:
# If a PyTorch model is to be loaded from a file
checkpoint = torch.load(args.model_path)
model.load_state_dict(checkpoint['model_state_dict'])
start_epoch = 0
start_iteration = 0
if resume:
# Resume training from last saved checkpoint
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['model_state_dict'])
start_epoch = checkpoint['epoch']
start_iteration = checkpoint['iteration']
else:
pass
# Loss - cross entropy between predicted scores (unnormalized) and class labels
# http://pytorch.org/docs/master/nn.html?highlight=crossentropyloss#crossentropyloss
criterion = nn.CrossEntropyLoss()
if cuda:
model = model.cuda()
criterion = criterion.cuda()
# -----------------------------------------------------------------------------
# 3. Optimizer
# -----------------------------------------------------------------------------
params = filter(lambda p: p.requires_grad, model.parameters())
# parameters with p.requires_grad=False are not updated during training
# this can be specified when defining the nn.Modules during model creation
if 'optim' in cfg.keys():
if cfg['optim'].lower() == 'sgd':
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
elif cfg['optim'].lower() == 'adam':
optim = torch.optim.Adam(params,
lr=cfg['lr'],
weight_decay=cfg['weight_decay'])
else:
raise NotImplementedError('Optimizers: SGD or Adam')
else:
optim = torch.optim.SGD(params,
lr=cfg['lr'],
momentum=cfg['momentum'],
weight_decay=cfg['weight_decay'])
if resume:
optim.load_state_dict(checkpoint['optim_state_dict'])
# -----------------------------------------------------------------------------
# [optional] Sanity-check: forward pass with a single batch
# -----------------------------------------------------------------------------
DEBUG = False
if DEBUG:
dataiter = iter(val_loader)
img, label = dataiter.next()
print 'Labels: ' + str(label.size()) # batchSize x num_class
print 'Input: ' + str(img.size()) # batchSize x 3 x 224 x224
im = img.squeeze().numpy()
im = im[0, :, :, :] # get first image in the batch
im = im.transpose((1, 2, 0)) # permute to 224x224x3
f = plt.figure()
plt.imshow(im)
plt.savefig(
'sanity-check-im.jpg') # save transformed image in current folder
inputs = Variable(img)
if cuda:
inputs = inputs.cuda()
model.eval()
outputs = model(inputs)
print 'Network output: ' + str(outputs.size())
model.train()
else:
pass
# -----------------------------------------------------------------------------
# Training
# -----------------------------------------------------------------------------
trainer = train.Trainer(
cuda=cuda,
model=model,
criterion=criterion,
optimizer=optim,
train_loader=train_loader,
val_loader=val_loader,
out=out,
max_iter=cfg['max_iteration'],
interval_validate=cfg.get('interval_validate', len(train_loader)),
)
trainer.epoch = start_epoch
trainer.iteration = start_iteration
trainer.train()