def main():
global args, best_prec1
args = parser.parse_args()
with open(args.config) as f:
config = yaml.load(f)
for key in config:
for k, v in config[key].items():
setattr(args, k, v)
print('Enabled distributed training.')
rank, world_size = init_dist(backend='nccl', port=args.port)
args.rank = rank
args.world_size = world_size
np.random.seed(args.seed * args.rank)
torch.manual_seed(args.seed * args.rank)
torch.cuda.manual_seed(args.seed * args.rank)
torch.cuda.manual_seed_all(args.seed * args.rank)
# create model
print("=> creating model '{}'".format(args.model))
if args.SinglePath:
architecture = args.arch
scale_list = 8 * [1.0]
scale_ids = [
6, 5, 3, 5, 2, 6, 3, 4, 2, 5, 7, 5, 4, 6, 7, 4, 4, 5, 4, 3
]
channels_scales = []
for i in range(len(scale_ids)):
channels_scales.append(scale_list[scale_ids[i]])
model = ShuffleNetV2_OneShot(args=args,
architecture=architecture,
channels_scales=channels_scales)
model.cuda()
broadcast_params(model)
# auto resume from a checkpoint
remark = 'imagenet_'
if args.remark != 'none':
remark += args.remark
args.save = 'search-{}-{}-{}'.format(args.save,
time.strftime("%Y%m%d-%H%M%S"),
remark)
args.save_log = 'nas-{}-{}'.format(time.strftime("%Y%m%d-%H%M%S"), remark)
generate_date = str(datetime.now().date())
path = os.path.join(generate_date, args.save)
if args.rank == 0:
log_format = '%(asctime)s %(message)s'
utils.create_exp_dir(generate_date,
path,
scripts_to_save=glob.glob('*.py'))
logging.basicConfig(stream=sys.stdout,
level=logging.INFO,
format=log_format,
datefmt='%m/%d %I:%M:%S %p')
fh = logging.FileHandler(os.path.join(path, 'log.txt'))
fh.setFormatter(logging.Formatter(log_format))
logging.getLogger().addHandler(fh)
logging.info("args = %s", args)
writer = SummaryWriter('./runs/' + generate_date + '/' + args.save_log)
else:
writer = None
#model_dir = args.model_dir
model_dir = path
start_epoch = 0
if args.evaluate:
load_state_ckpt(args.checkpoint_path, model)
cudnn.benchmark = True
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
val_dataset = ImagenetDataset(
args.val_root, args.val_source,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
val_sampler = DistributedSampler(val_dataset)
val_loader = DataLoader(val_dataset,
batch_size=50,
shuffle=False,
num_workers=args.workers,
pin_memory=False,
sampler=val_sampler)
if args.evaluate:
validate(val_loader, model, 0, writer, logging)
return
def main():
global args, best_prec1
args = parser.parse_args()
with open(args.config) as f:
config = yaml.load(f)
for key in config:
for k, v in config[key].items():
setattr(args, k, v)
print('Enabled distributed training.')
rank, world_size = init_dist(backend='nccl', port=args.port)
args.rank = rank
args.world_size = world_size
# create model
print("=> creating model '{}'".format(args.model))
if 'resnetv1sn' in args.model:
model = models.__dict__[args.model](
using_moving_average=args.using_moving_average,
using_bn=args.using_bn,
last_gamma=args.last_gamma)
else:
model = models.__dict__[args.model](
using_moving_average=args.using_moving_average,
using_bn=args.using_bn)
model.cuda()
broadcast_params(model)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.base_lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# auto resume from a checkpoint
model_dir = args.model_dir
start_epoch = 0
if args.rank == 0 and not os.path.exists(model_dir):
os.makedirs(model_dir)
if args.evaluate:
load_state_ckpt(args.checkpoint_path, model)
else:
best_prec1, start_epoch = load_state(model_dir,
model,
optimizer=optimizer)
if args.rank == 0:
writer = SummaryWriter(model_dir)
else:
writer = None
cudnn.benchmark = True
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = ImagenetDataset(
args.train_root, args.train_source,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
ColorAugmentation(),
normalize,
]))
val_dataset = ImagenetDataset(
args.val_root, args.val_source,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_sampler = DistributedSampler(train_dataset)
val_sampler = DistributedSampler(val_dataset)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size // args.world_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False,
sampler=train_sampler)
val_loader = DataLoader(val_dataset,
batch_size=args.batch_size // args.world_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False,
sampler=val_sampler)
if args.evaluate:
validate(val_loader, model, criterion, 0, writer)
return
niters = len(train_loader)
lr_scheduler = LRScheduler(optimizer, niters, args)
for epoch in range(start_epoch, args.epochs):
train_sampler.set_epoch(epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, lr_scheduler, epoch,
writer)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch, writer)
if rank == 0:
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
model_dir, {
'epoch': epoch + 1,
'model': args.model,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)
def main():
global args, best_prec1
args = parser.parse_args()
with open(args.config) as f:
config = yaml.load(f)
for key in config:
for k, v in config[key].items():
setattr(args, k, v)
print('Enabled distributed training.')
rank, world_size = init_dist(
backend='nccl', port=args.port)
args.rank = rank
args.world_size = world_size
np.random.seed(args.seed*args.rank)
torch.manual_seed(args.seed*args.rank)
torch.cuda.manual_seed(args.seed*args.rank)
torch.cuda.manual_seed_all(args.seed*args.rank)
# create model
print("=> creating model '{}'".format(args.model))
if args.SinglePath:
architecture = 20*[0]
channels_scales = 20*[1.0]
#load derived child network
log_alpha = torch.load(args.checkpoint_path, map_location='cuda:{}'.format(torch.cuda.current_device()))['state_dict']['log_alpha']
weights = torch.zeros_like(log_alpha).scatter_(1, torch.argmax(log_alpha, dim = -1).view(-1,1), 1)
model = ShuffleNetV2_OneShot(args=args, architecture=architecture, channels_scales=channels_scales, weights=weights)
model.cuda()
broadcast_params(model)
for v in model.parameters():
if v.requires_grad:
if v.grad is None:
v.grad = torch.zeros_like(v)
model.log_alpha.grad = torch.zeros_like(model.log_alpha)
if not args.retrain:
load_state_ckpt(args.checkpoint_path, model)
checkpoint = torch.load(args.checkpoint_path, map_location='cuda:{}'.format(torch.cuda.current_device()))
args.base_lr = checkpoint['optimizer']['param_groups'][0]['lr']
if args.reset_bn_stat:
model._reset_bn_running_stats()
# define loss function (criterion) and optimizer
criterion = CrossEntropyLoss(smooth_eps=0.1, smooth_dist=(torch.ones(1000)*0.001).cuda()).cuda()
wo_wd_params = []
wo_wd_param_names = []
network_params = []
network_param_names = []
for name, mod in model.named_modules():
#if isinstance(mod, (nn.BatchNorm2d, SwitchNorm2d)):
if isinstance(mod, nn.BatchNorm2d):
for key, value in mod.named_parameters():
wo_wd_param_names.append(name+'.'+key)
for key, value in model.named_parameters():
if key != 'log_alpha':
if value.requires_grad:
if key in wo_wd_param_names:
wo_wd_params.append(value)
else:
network_params.append(value)
network_param_names.append(key)
params = [
{'params': network_params,
'lr': args.base_lr,
'weight_decay': args.weight_decay },
{'params': wo_wd_params,
'lr': args.base_lr,
'weight_decay': 0.},
]
param_names = [network_param_names, wo_wd_param_names]
if args.rank == 0:
print('>>> params w/o weight decay: ', wo_wd_param_names)
optimizer = torch.optim.SGD(params, momentum=args.momentum)
arch_optimizer=None
# auto resume from a checkpoint
remark = 'imagenet_'
remark += 'epo_' + str(args.epochs) + '_layer_' + str(args.layers) + '_batch_' + str(args.batch_size) + '_lr_' + str(float("{0:.2f}".format(args.base_ lr))) + '_seed_' + str(args.seed)
if args.remark != 'none':
remark += '_'+args.remark
args.save = 'search-{}-{}-{}'.format(args.save, time.strftime("%Y%m%d-%H%M%S"), remark)
args.save_log = 'nas-{}-{}'.format(time.strftime("%Y%m%d-%H%M%S"), remark)
generate_date = str(datetime.now().date())
path = os.path.join(generate_date, args.save)
if args.rank == 0:
log_format = '%(asctime)s %(message)s'
utils.create_exp_dir(generate_date, path, scripts_to_save=glob.glob('*.py'))
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
format=log_format, datefmt='%m/%d %I:%M:%S %p')
fh = logging.FileHandler(os.path.join(path, 'log.txt'))
fh.setFormatter(logging.Formatter(log_format))
logging.getLogger().addHandler(fh)
logging.info("args = %s", args)
writer = SummaryWriter('./runs/' + generate_date + '/' + args.save_log)
else:
writer = None
#model_dir = args.model_dir
model_dir = path
start_epoch = 0
if args.evaluate:
load_state_ckpt(args.checkpoint_path, model)
else:
best_prec1, start_epoch = load_state(model_dir, model, optimizer=optimizer)
cudnn.benchmark = True
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = ImagenetDataset(
args.train_root,
args.train_source,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
train_dataset_wo_ms = ImagenetDataset(
args.train_root,
args.train_source,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_dataset = ImagenetDataset(
args.val_root,
args.val_source,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_sampler = DistributedSampler(train_dataset)
val_sampler = DistributedSampler(val_dataset)
train_loader = DataLoader(
train_dataset, batch_size=args.batch_size//args.world_size, shuffle=False,
num_workers=args.workers, pin_memory=False, sampler=train_sampler)
train_loader_wo_ms = DataLoader(
train_dataset_wo_ms, batch_size=args.batch_size//args.world_size, shuffle=False,
num_workers=args.workers, pin_memory=False, sampler=train_sampler)
val_loader = DataLoader(
val_dataset, batch_size=50, shuffle=False,
num_workers=args.workers, pin_memory=False, sampler=val_sampler)
if args.evaluate:
validate(val_loader, model, criterion, 0, writer, logging)
return
niters = len(train_loader)
lr_scheduler = LRScheduler(optimizer, niters, args)
for epoch in range(start_epoch, args.epochs):
train_sampler.set_epoch(epoch)
if args.rank == 0 and args.SinglePath:
logging.info('epoch %d', epoch)
# evaluate on validation set after loading the model
if epoch == 0 and not args.reset_bn_stat:
prec1 = validate(val_loader, model, criterion, epoch, writer, logging)
# train for one epoch
if epoch >= args.epochs - 5 and args.lr_mode == 'step' and args.off_ms and args.retrain:
train(train_loader_wo_ms, model, criterion, optimizer, arch_optimizer, lr_scheduler, epoch, writer, logging)
else:
train(train_loader, model, criterion, optimizer, arch_optimizer, lr_scheduler, epoch, writer, logging)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch, writer, logging)
if rank == 0:
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(model_dir, {
'epoch': epoch + 1,
'model': args.model,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)
def main():
global args, best_prec1
args = parser.parse_args()
with open(args.config) as f:
config = yaml.load(f)
for key in config:
for k, v in config[key].items():
setattr(args, k, v)
print('Enabled distributed training.')
rank, world_size = init_dist(
backend='nccl', port=args.port)
args.rank = rank
args.world_size = world_size
np.random.seed(args.seed*args.rank)
torch.manual_seed(args.seed*args.rank)
torch.cuda.manual_seed(args.seed*args.rank)
torch.cuda.manual_seed_all(args.seed*args.rank)
print('random seed: ', args.seed*args.rank)
# create model
print("=> creating model '{}'".format(args.model))
if args.SinglePath:
architecture = 20*[0]
channels_scales = 20*[1.0]
model = ShuffleNetV2_OneShot(args=args, architecture=architecture, channels_scales=channels_scales)
model.cuda()
broadcast_params(model)
for v in model.parameters():
if v.requires_grad:
if v.grad is None:
v.grad = torch.zeros_like(v)
model.log_alpha.grad = torch.zeros_like(model.log_alpha)
criterion = CrossEntropyLoss(smooth_eps=0.1, smooth_dist=(torch.ones(1000)*0.001).cuda()).cuda()
wo_wd_params = []
wo_wd_param_names = []
network_params = []
network_param_names = []
for name, mod in model.named_modules():
if isinstance(mod, nn.BatchNorm2d):
for key, value in mod.named_parameters():
wo_wd_param_names.append(name+'.'+key)
for key, value in model.named_parameters():
if key != 'log_alpha':
if value.requires_grad:
if key in wo_wd_param_names:
wo_wd_params.append(value)
else:
network_params.append(value)
network_param_names.append(key)
params = [
{'params': network_params,
'lr': args.base_lr,
'weight_decay': args.weight_decay },
{'params': wo_wd_params,
'lr': args.base_lr,
'weight_decay': 0.},
]
param_names = [network_param_names, wo_wd_param_names]
if args.rank == 0:
print('>>> params w/o weight decay: ', wo_wd_param_names)
optimizer = torch.optim.SGD(params, momentum=args.momentum)
if args.SinglePath:
arch_optimizer = torch.optim.Adam(
[param for name, param in model.named_parameters() if name == 'log_alpha'],
lr=args.arch_learning_rate,
betas=(0.5, 0.999),
weight_decay=args.arch_weight_decay
)
# auto resume from a checkpoint
remark = 'imagenet_'
remark += 'epo_' + str(args.epochs) + '_layer_' + str(args.layers) + '_batch_' + str(args.batch_size) + '_lr_' + str(args.base_lr) + '_seed_' + str(args.seed) + '_pretrain_' + str(args.pretrain_epoch)
if args.early_fix_arch:
remark += '_early_fix_arch'
if args.flops_loss:
remark += '_flops_loss_' + str(args.flops_loss_coef)
if args.remark != 'none':
remark += '_'+args.remark
args.save = 'search-{}-{}-{}'.format(args.save, time.strftime("%Y%m%d-%H%M%S"), remark)
args.save_log = 'nas-{}-{}'.format(time.strftime("%Y%m%d-%H%M%S"), remark)
generate_date = str(datetime.now().date())
path = os.path.join(generate_date, args.save)
if args.rank == 0:
log_format = '%(asctime)s %(message)s'
utils.create_exp_dir(generate_date, path, scripts_to_save=glob.glob('*.py'))
logging.basicConfig(stream=sys.stdout, level=logging.INFO,
format=log_format, datefmt='%m/%d %I:%M:%S %p')
fh = logging.FileHandler(os.path.join(path, 'log.txt'))
fh.setFormatter(logging.Formatter(log_format))
logging.getLogger().addHandler(fh)
logging.info("args = %s", args)
writer = SummaryWriter('./runs/' + generate_date + '/' + args.save_log)
else:
writer = None
model_dir = path
start_epoch = 0
if args.evaluate:
load_state_ckpt(args.checkpoint_path, model)
else:
best_prec1, start_epoch = load_state(model_dir, model, optimizer=optimizer)
cudnn.benchmark = True
cudnn.enabled = True
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = ImagenetDataset(
args.train_root,
args.train_source,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
train_dataset_wo_ms = ImagenetDataset(
args.train_root,
args.train_source,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_dataset = ImagenetDataset(
args.val_root,
args.val_source,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_sampler = DistributedSampler(train_dataset)
val_sampler = DistributedSampler(val_dataset)
train_loader = DataLoader(
train_dataset, batch_size=args.batch_size//args.world_size, shuffle=False,
num_workers=args.workers, pin_memory=False, sampler=train_sampler)
train_loader_wo_ms = DataLoader(
train_dataset_wo_ms, batch_size=args.batch_size//args.world_size, shuffle=False,
num_workers=args.workers, pin_memory=False, sampler=train_sampler)
val_loader = DataLoader(
val_dataset, batch_size=50, shuffle=False,
num_workers=args.workers, pin_memory=False, sampler=val_sampler)
if args.evaluate:
validate(val_loader, model, criterion, 0, writer, logging)
return
niters = len(train_loader)
lr_scheduler = LRScheduler(optimizer, niters, args)
for epoch in range(start_epoch, 85):
train_sampler.set_epoch(epoch)
if args.early_fix_arch:
if len(model.fix_arch_index.keys()) > 0:
for key, value_lst in model.fix_arch_index.items():
model.log_alpha.data[key, :] = value_lst[1]
sort_log_alpha = torch.topk(F.softmax(model.log_alpha.data, dim=-1), 2)
argmax_index = (sort_log_alpha[0][:,0] - sort_log_alpha[0][:,1] >= 0.3)
for id in range(argmax_index.size(0)):
if argmax_index[id] == 1 and id not in model.fix_arch_index.keys():
model.fix_arch_index[id] = [sort_log_alpha[1][id,0].item(), model.log_alpha.detach().clone()[id, :]]
if args.rank == 0 and args.SinglePath:
logging.info('epoch %d', epoch)
logging.info(model.log_alpha)
logging.info(F.softmax(model.log_alpha, dim=-1))
logging.info('flops %fM', model.cal_flops())
# train for one epoch
if epoch >= args.epochs - 5 and args.lr_mode == 'step' and args.off_ms:
train(train_loader_wo_ms, model, criterion, optimizer, arch_optimizer, lr_scheduler, epoch, writer, logging)
else:
train(train_loader, model, criterion, optimizer, arch_optimizer, lr_scheduler, epoch, writer, logging)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch, writer, logging)
if args.gen_max_child:
args.gen_max_child_flag = True
prec1 = validate(val_loader, model, criterion, epoch, writer, logging)
args.gen_max_child_flag = False
if rank == 0:
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(model_dir, {
'epoch': epoch + 1,
'model': args.model,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)
def main(args):
print(args)
model = ArcFaceWithLoss(args.backbone, 85742, args.norm_func,
args.embedding_size, args.use_se)
model.cuda()
load_state_ckpt(args.checkpoint_path, model)
facescrub_out = os.path.join('./' + args.algo + '_' + args.output,
'facescrub')
megaface_out = os.path.join('./' + args.algo + '_' + args.output,
'megaface')
i = 0
succ = 0
buffer = []
for line in open(args.facescrub_lst, 'r'):
if i % 1000 == 0:
print("writing fs", i, succ)
i += 1
image_path = line.strip()
_path = image_path.split('/')
a, b = _path[-2], _path[-1]
out_dir = os.path.join(facescrub_out, a)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
image_path = os.path.join(args.facescrub_root, image_path)
img = read_img(image_path)
if img is None:
print('read error:', image_path)
continue
out_path = os.path.join(out_dir, b + "_%s.bin" % (args.algo))
item = (img, out_path)
buffer.append(item)
if len(buffer) == args.batch_size:
get_and_write(buffer, model)
buffer = []
succ += 1
if len(buffer) > 0:
get_and_write(buffer, model)
buffer = []
print('fs stat', i, succ)
i = 0
succ = 0
buffer = []
for line in open(args.megaface_lst, 'r'):
if i % 1000 == 0:
print("writing mf", i, succ)
i += 1
image_path = line.strip()
_path = image_path.split('/')
a1, a2, b = _path[-3], _path[-2], _path[-1]
out_dir = os.path.join(megaface_out, a1, a2)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
# continue
# print(landmark)
image_path = os.path.join(args.megaface_root, image_path)
img = read_img(image_path)
if img is None:
print('read error:', image_path)
continue
out_path = os.path.join(out_dir, b + "_%s.bin" % (args.algo))
item = (img, out_path)
buffer.append(item)
if len(buffer) == args.batch_size:
get_and_write(buffer, model)
buffer = []
succ += 1
if len(buffer) > 0:
get_and_write(buffer, model)
buffer = []
print('mf stat', i, succ)
def main_worker(gpu, ngpus_per_node, args):
args.gpu = gpu
args.rank = args.rank * ngpus_per_node + gpu
if args.distributed:
print(args.backend, args.world_size, args.rank)
dist.init_process_group(backend=args.backend,
init_method='tcp://127.0.0.1:6668',
world_size=args.world_size,
rank=args.rank)
print('Enabled distributed training.')
# create model
print("=> creating model '{}'".format(args.model))
model = models.__dict__[args.model](N=args.N, M=args.M)
torch.cuda.set_device(args.gpu)
ipClass = PruningMethodTransposableBlockL1(block_size=args.M, topk=args.N)
if args.load_mask:
load_state_and_masks(model, args)
print("Masks loaded!")
else:
for n, m in model.named_modules():
if isinstance(m, SparseConvTranspose) or isinstance(
m, SparseLinearTranspose):
# m.maskBuff.data = ipClass.compute_mask(m.weight, torch.ones_like(m.weight))
setattr(
m.weight, "mask",
ipClass.compute_mask(m.weight, torch.ones_like(m.weight)))
if args.save_mask:
save_masks(model, args)
print("Masks saved!")
model.cuda(args.gpu)
#args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
#broadcast_params(model)
print(model)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
if args.sparse_optimizer:
optimizer = sparse_optimizer.SGD(model.parameters(),
args.base_lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.SGD(model.parameters(),
args.base_lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# auto resume from a checkpoint
model_dir = args.model_dir
start_epoch = 0
best_prec1 = 0
if args.rank == 0 and not os.path.exists(model_dir):
os.makedirs(model_dir)
if args.evaluate:
load_state_ckpt(args.checkpoint_path, model)
else:
best_prec1, start_epoch = load_state(model_dir,
model,
optimizer=optimizer)
if args.rank == 0 or not args.distributed:
writer = SummaryWriter(model_dir)
else:
writer = None
cudnn.benchmark = True
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
args.train_root,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
ColorAugmentation(),
normalize,
]))
val_dataset = datasets.ImageFolder(
args.val_root,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = DistributedSampler(train_dataset)
val_sampler = DistributedSampler(val_dataset)
else:
train_sampler = None
val_sampler = None
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size // args.world_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False,
sampler=train_sampler)
val_loader = DataLoader(val_dataset,
batch_size=args.batch_size // args.world_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False,
sampler=val_sampler)
if args.evaluate:
validate(val_loader, model, criterion, 0, writer)
return
niters = len(train_loader)
lr_scheduler = LRScheduler(optimizer, niters, args)
for epoch in range(start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, lr_scheduler, epoch,
writer, args)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion, epoch, writer, args)
if args.rank == 0:
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
model_dir, {
'epoch': epoch + 1,
'model': args.model,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)