def main():
args = parser.parse_args()
assert torch.cuda.is_available()
image_ph = tf.placeholder(tf.uint8, (None, None, 3))
image_proc = preprocess_for_eval(image_ph, args.image_size,
args.image_size)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
model = NetworkImageNet(args.num_conv_filters, args.num_classes,
args.num_cells, False, PNASNet)
model.drop_path_prob = 0
model.eval()
model.load_state_dict(torch.load('data/PNASNet-5_Large.pth'))
model = model.cuda()
c1, c5 = 0, 0
val_dataset = datasets.ImageFolder(args.valdir)
for i, (image, label) in enumerate(val_dataset):
tf_image_proc = sess.run(image_proc, feed_dict={image_ph: image})
image = torch.from_numpy(tf_image_proc.transpose((2, 0, 1)))
image = Variable(image).cuda()
logits, _ = model(image.unsqueeze(0))
top5 = logits.data.cpu().numpy().squeeze().argsort()[::-1][:5]
top1 = top5[0]
if label + 1 == top1:
c1 += 1
if label + 1 in top5:
c5 += 1
print('Test: [{0}/{1}]\t'
'Prec@1 {2:.3f}\t'
'Prec@5 {3:.3f}\t'.format(i + 1, len(val_dataset), c1 / (i + 1.),
c5 / (i + 1.)))
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
if args.parallel:
model = nn.DataParallel(model).cuda()
else:
model = model.cuda()
model.load_state_dict(
torch.load(args.model_path, map_location='cuda:0')['state_dict'])
print("param size = {:.1f}MB".format(
floor(utils.count_parameters_in_MB(model), 1)))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
validdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=4)
model.drop_path_prob = args.drop_path_prob
input = torch.randn(1, 3, 224, 224)
input = input.cuda()
valid_acc_top1, valid_acc_top5, valid_obj = infer(valid_queue, model,
criterion)
logging.info('valid_acc_top1 %f', valid_acc_top1)
logging.info('valid_acc_top5 %f', valid_acc_top5)
def get_imagenet_tuned_model(load_weights=True):
network = NetworkImageNet(48, IMAGENET_CLASSES, 14, True, 0.4,
IMAGENET_TUNED,
CrossEntropyLabelSmooth(IMAGENET_CLASSES, 0.1))
if load_weights:
device = torch.device('cpu')
state_dict = torch.load('weights/imagenet_tuned.pt',
map_location=device)
# state_dict = {k:v for k,v in state_dict.items() if not 'total_ops' in k and not 'total_params' in k}
network.load_state_dict(state_dict)
return network
def eval_arch(genotype_file, ckpt_path):
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled=True
torch.cuda.manual_seed(args.seed)
logging.info("args = %s", args)
# tmp_dict = json.load(open(genotype_file,'r'))
if args.arch is not None:
genotype = eval("genotypes.%s" % args.arch)
# genotype = genotypes.Genotype(**tmp_dict)
print(genotype)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary, genotype)
model.drop_path_prob = 0.
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model).cuda()
else:
model = model.cuda()
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
model.load_state_dict(torch.load(ckpt_path), strict=False)
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
validdir = os.path.join(args.data, 'ILSVRC2012_img_val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
valid_queue = torch.utils.data.DataLoader(
valid_data, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=4)
valid_acc_top1, valid_acc_top5, valid_obj = infer(valid_queue, model, criterion)
logging.info('valid_acc_top1 %f', valid_acc_top1)
logging.info('valid_acc_top5 %f', valid_acc_top5)
def main():
np.random.seed(args.seed)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info(
"device = %s" %
'cuda:{}'.format(args.gpu) if torch.cuda.is_available() else 'cpu')
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype).to(device)
model.load_state_dict(torch.load(args.model_path)["state_dict"])
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss().to(device)
validdir = os.path.join(args.data, "val")
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]),
)
valid_queue = torch.utils.data.DataLoader(
valid_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=4,
)
model.drop_path_prob = args.drop_path_prob
valid_acc_top1, valid_acc_top5, valid_obj = infer(valid_queue, model,
criterion)
logging.info("valid_acc_top1 %f", valid_acc_top1)
logging.info("valid_acc_top5 %f", valid_acc_top5)
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
model = model.cuda()
model.load_state_dict(torch.load(args.model_path)['state_dict'])
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
validdir = os.path.join(args.data, 'valid')
normalize = transforms.Normalize(mean=[0.4802, 0.4481, 0.3975],
std=[0.2302, 0.2265, 0.2262])
valid_data = dset.ImageFolder(
validdir, transforms.Compose([
transforms.ToTensor(),
normalize,
]))
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=4)
model.drop_path_prob = args.drop_path_prob
valid_acc_top1, valid_acc_top5, valid_obj = infer(valid_queue, model,
criterion)
logging.info('valid_acc_top1 %f', valid_acc_top1)
logging.info('valid_acc_top5 %f', valid_acc_top5)
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
cudnn.enabled = True
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
model = nn.DataParallel(model)
model = model.cuda()
model.load_state_dict(torch.load(args.model_path)['state_dict'],
strict=False)
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
validdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=False,
num_workers=4)
model.module.drop_path_prob = 0.0
valid_acc_top1, valid_acc_top5, valid_obj = infer(valid_queue, model,
criterion)
logging.info('Valid_acc_top1 %f', valid_acc_top1)
logging.info('Valid_acc_top5 %f', valid_acc_top5)
def main():
if not torch.cuda.is_available():
print('No GPU device available')
sys.exit(1)
## step 1 construct the selected network
genotype = eval("genotypes.%s" % args.selected_arch)
CLASSES = 1000
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.cuda()
else:
model = model.cuda()
## step 2 load pretrained model parameter
model_CKPT = torch.load(args.model_path)
model.load_state_dict(model_CKPT['state_dict'])
model.module.drop_path_prob = 0
model.drop_path_prob = 0
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
## step 3 load test data
valid_queue = load_data_cifar(args)
## step 4. inference on test data
valid_acc_top1, valid_acc_top5, valid_obj = infer(valid_queue, model,
criterion)
print('-----------------------------------------------')
print('Valid_acc_top1: %f, Valid_acc_top5: %f' %
(valid_acc_top1, valid_acc_top5))
print('-----------------------------------------------')
def main():
if not torch.cuda.is_available():
logging.info('No GPU device available')
sys.exit(1)
np.random.seed(args.seed)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled=True
torch.cuda.manual_seed(args.seed)
logging.info("args = %s", args)
logging.info("unparsed_args = %s", unparsed)
num_gpus = torch.cuda.device_count()
if not os.path.exists(args.base_path):
os.makedirs(args.base_path)
ckpt_dir = os.path.join(args.base_path, "ImageNet")
if not os.path.exists(ckpt_dir):
os.makedirs(ckpt_dir)
# ckpt_dir_1 = os.path.join(args.base_path, "ImageNet-1")
# if not os.path.exists(ckpt_dir):
# os.makedirs(ckpt_dir_1)
# print(arhs.arch)
if args.arch is not None:
genotype = eval("genotypes.%s" % args.arch)
elif args.base_path is not None and args.genotype_name is not None:
genotype_path = os.path.join(args.base_path, 'results_of_7q/genotype')
genotype = get_genotype(genotype_path, args.genotype_name)
else:
raise(ValueError("the parser input arch, genotype_path, genotype_name should not be all None"))
# print(genotype)
# genotype = eval("genotypes.%s" % args.arch)
print('---------Genotype---------')
logging.info(genotype)
print('--------------------------')
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary, genotype)
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.cuda()
else:
model = model.cuda()
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
optimizer = torch.optim.SGD(
model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay
)
# data_dir = os.path.join(args.tmp_data_dir, 'imagenet')
# traindir = os.path.join(data_dir, 'train')
# validdir = os.path.join(data_dir, 'val')
traindir = os.path.join(args.data, 'ILSVRC2012_img_train')
validdir = os.path.join(args.data, 'ILSVRC2012_img_val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_data = dset.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(
brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.2),
transforms.ToTensor(),
normalize,
]))
train_queue = torch.utils.data.DataLoader(
train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=args.workers)
valid_queue = torch.utils.data.DataLoader(
valid_data, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=args.workers)
best_acc_top1 = 0
best_acc_top5 = 0
ckpt_file = None
start_epoch = 0
if not args.from_scratch:
# check wheter have pre-trained models
if args.load_file is not None and os.path.exists(args.load_file):
ckpt_file = args.load_file
else:
# deal with negative names
files = os.listdir(ckpt_dir)
for f in files:
tmp = f.split('.')
if len(tmp) > 2: continue
tmp = int(tmp[0].split('_')[1])
if tmp > start_epoch:
start_epoch = tmp
ckpt_file = os.path.join(ckpt_dir, f)
if ckpt_file is not None:
logging.info('====== Load ckpt ======')
logging.info("Loading from %s"%ckpt_file)
start_epoch = 249
# if num_gpus > 1:
# model.module.load_state_dict(checkpoint['state_dict'])
# else:
# model.load_state_dict(checkpoint['state_dict'])
# start_epoch = int(checkpoint['epoch']) + 1
# optimizer.load_state_dict(checkpoint['optimizer'])
# best_acc_top1 = float(checkpoint['best_acc_top1'])
# logging.info("Training Start at %d"%start_epoch)
logging.info("Training Start at %d"%start_epoch)
sat= torch.load(ckpt_file)
# print(type(sat))
# if num_gpus > 1:
# model.load_state_dict(sat)
# else:
# model = model.cuda()
# new_stat={}
# for k,v in sat.items():
# # "module.stem0.0.weight"
# new_k = k#[7:]
# new_stat[new_k] =copy.deepcopy(v)
model.load_state_dict(sat)
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.decay_period, gamma=args.gamma, last_epoch=start_epoch)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, float(args.epochs))
for epoch in range(start_epoch, args.epochs):
if args.lr_scheduler == 'cosine':
scheduler.step()
current_lr = scheduler.get_lr()[0]
elif args.lr_scheduler == 'linear':
current_lr = adjust_lr(optimizer, epoch)
else:
print('Wrong lr type, exit')
sys.exit(1)
logging.info('Epoch: %d lr %e', epoch, current_lr)
if epoch < 5 and args.batch_size > 256:
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr * (epoch + 1) / 5.0
logging.info('Warming-up Epoch: %d, LR: %e', epoch, current_lr * (epoch + 1) / 5.0)
if num_gpus > 1:
model.module.drop_path_prob = args.drop_path_prob * epoch / args.epochs
else:
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
epoch_start = time.time()
train_acc, train_obj = train(train_queue, model, criterion_smooth, optimizer)
logging.info('Train_acc: %f', train_acc)
valid_acc_top1, valid_acc_top5, valid_obj = infer(valid_queue, model, criterion)
logging.info('Valid_acc_top1: %f', valid_acc_top1)
logging.info('Valid_acc_top5: %f', valid_acc_top5)
epoch_duration = time.time() - epoch_start
logging.info('Epoch time: %ds.', epoch_duration)
# save current epoch model, and remove previous model
try:
last_model = os.path.join(ckpt_dir, 'weights_%d.pt'%(epoch-1))
os.remove(last_model)
except:
pass
ckpt = {
'epoch': epoch,
'state_dict': model.state_dict(),
'best_acc_top1': best_acc_top1,
'optimizer' : optimizer.state_dict(),
}
utils.save(model, os.path.join(ckpt_dir, 'weights_%d.pt'%(epoch)))
if valid_acc_top1 > best_acc_top1:
try:
last_model = os.path.join(ckpt_dir, 'weights_%.3f.pt'%(best_acc_top1))
os.remove(last_model)
except:
pass
ckpt = {
'epoch': epoch,
'state_dict': model.state_dict(),
'best_acc_top1': best_acc_top1,
'optimizer' : optimizer.state_dict(),
}
utils.save(model, os.path.join(ckpt_dir, 'weights_%.3f.pt'%(valid_acc_top1)))
best_acc_top1 = valid_acc_top1
best_acc_top5 = valid_acc_top5
# is_best = False
# if valid_acc_top5 > best_acc_top5:
# best_acc_top5 = valid_acc_top5
# if valid_acc_top1 > best_acc_top1:
# best_acc_top1 = valid_acc_top1
# is_best = True
# utils.save_checkpoint({
# 'epoch': epoch + 1,
# 'state_dict': model.state_dict(),
# 'best_acc_top1': best_acc_top1,
# 'optimizer' : optimizer.state_dict(),
# }, is_best, args.save)
print("the best acc: %f(Top1); %f(Top5)"%(best_acc_top1, best_acc_top5))
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
rank, world_size = dist_util.dist_init(args.port, 'nccl')
np.random.seed(args.seed)
# torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
# logging.info('gpu device = %d' % args.gpu)
if rank == 0:
generate_date = str(datetime.now().date())
utils.create_exp_dir(generate_date, args.save, 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(args.save, 'log.txt'))
fh.setFormatter(logging.Formatter(log_format))
logging.getLogger().addHandler(fh)
logging.info("args = %s", args)
logger = tensorboardX.SummaryWriter('./runs/eval_imagenet_{}_{}'.format(args.arch, args.remark))
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary, genotype)
model = model.cuda()
if rank == 0:
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
if args.warm_start:
lr = args.warm_start_lr / args.warm_start_gamma
else:
lr = args.learning_rate
optimizer = torch.optim.SGD(
model.parameters(),
lr,
momentum=args.momentum,
weight_decay=args.weight_decay
)
traindir = os.path.join(args.data, 'train')
validdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(
brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.2),
transforms.ToTensor(),
normalize,
])
train_dataset = get_dataset(traindir, os.path.join(args.data, 'meta/train.txt'), train_transform)
valid_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])
valid_dataset = get_dataset(validdir, os.path.join(args.data, 'meta/val.txt'), valid_transform)
# train_queue = torch.utils.data.DataLoader(
# train_dataset, batch_size=args.batch_size // world_size, sampler=DistributedSampler(train_dataset),
# pin_memory=True, num_workers=4)
# valid_queue = torch.utils.data.DataLoader(
# valid_dataset, batch_size=args.batch_size // world_size, sampler=DistributedSampler(valid_dataset),
# pin_memory=True, num_workers=4)
train_queue = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=4)
valid_queue = torch.utils.data.DataLoader(
valid_dataset, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=4)
if args.warm_start:
scheduler = utils.WarmStart(optimizer, gamma=args.warm_start_gamma)
else:
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.decay_period, gamma=args.gamma)
best_acc_top1 = 0
for epoch in range(args.epochs):
'''if epoch == 0 or epoch == 1:
for param_group in optimizer.param_groups:
param_group['lr'] = args.warm_up_learning_rate
elif epoch == 2:
for param_group in optimizer.param_groups:
param_group['lr'] = args.learning_rate
else:
scheduler.step()'''
scheduler.step()
if rank == 0:
logging.info('epoch %d lr %e', epoch, scheduler.get_lr()[0])
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
train_acc, train_obj = train(train_queue, model, criterion_smooth, optimizer, rank)
if rank == 0:
logging.info('train_acc %f', train_acc)
logger.add_scalar("epoch_train_acc", train_acc, epoch)
logger.add_scalar("epoch_train_loss", train_obj, epoch)
valid_acc_top1, valid_acc_top5, valid_obj = infer(valid_queue, model, criterion, rank)
if rank == 0:
logging.info('valid_acc_top1 %f', valid_acc_top1)
logging.info('valid_acc_top5 %f', valid_acc_top5)
logger.add_scalar("epoch_valid_acc_top1", valid_acc_top1, epoch)
logger.add_scalar("epoch_valid_acc_top5", valid_acc_top5, epoch)
is_best = False
if valid_acc_top1 > best_acc_top1:
best_acc_top1 = valid_acc_top1
is_best = True
if args.warm_start:
# if not is_best and not scheduler.lr_const:
if True:
if rank == 0:
logging.info('warm start ended lr %e', scheduler.get_lr()[0])
logging.info("=> loading checkpoint '{}'".format(args.save))
# checkpoint = torch.load(os.path.join(args.save, 'model_best.pth.tar'))
checkpoint = torch.load(os.path.join(args.save, 'model_best.pth.tar'),
map_location=lambda storage, loc: storage)
best_acc_top1 = checkpoint['best_acc_top1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.lr_const = True
if rank == 0:
logging.info('return to last checkpoint')
del checkpoint # dereference seems crucial
torch.cuda.empty_cache()
# args.start_epoch = checkpoint['epoch']
'''best_acc_top1 = checkpoint['best_acc_top1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.lr_const = True
if rank == 0:
logging.info('return to last checkpoint')'''
if rank == 0:
utils.save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_acc_top1': best_acc_top1,
'optimizer': optimizer.state_dict(),
}, is_best, args.save)
def main():
if not torch.cuda.is_available():
logging.info('No GPU device available')
sys.exit(1)
np.random.seed(args.seed)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
torch.cuda.set_device(initGpu)
logging.info("args = %s", args)
logging.info("unparsed_args = %s", unparsed)
num_gpus = torch.cuda.device_count()
genotype = eval("genotypes.%s" % args.arch)
print('---------Genotype---------')
logging.info(genotype)
print('--------------------------')
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
#start_epochs = 0
if num_gpus > 1:
model = nn.DataParallel(model, device_ids)
model = model.cuda()
else:
model = model.cuda()
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
if (not args.resume_path == ''):
state = utils.load_checkpoint(args.resume_path)
#start_epochs = state[epoch]
model.load_state_dict(state['state_dict'])
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
data_dir = args.tmp_data_dir
traindir = os.path.join(data_dir, 'train')
validdir = os.path.join(data_dir, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_data = dset.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.2),
transforms.ToTensor(),
normalize,
]))
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=args.workers)
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=args.workers)
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.decay_period, gamma=args.gamma)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, float(args.epochs))
best_acc_top1 = 0
best_acc_top5 = 0
for epoch in range(args.epochs):
if args.lr_scheduler == 'cosine':
scheduler.step()
current_lr = scheduler.get_lr()[0]
elif args.lr_scheduler == 'linear':
current_lr = adjust_lr(optimizer, epoch)
else:
print('Wrong lr type, exit')
sys.exit(1)
logging.info('Epoch: %d lr %e', epoch, current_lr)
if epoch < 5 and args.batch_size > 256:
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr * (epoch + 1) / 5.0
logging.info('Warming-up Epoch: %d, LR: %e', epoch,
current_lr * (epoch + 1) / 5.0)
if num_gpus > 1:
model.module.drop_path_prob = args.drop_path_prob * epoch / args.epochs
else:
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
epoch_start = time.time()
train_acc, train_obj = train(train_queue, model, criterion_smooth,
optimizer)
logging.info('Train_acc: %f', train_acc)
valid_acc_top1, valid_acc_top5, valid_obj = infer(
valid_queue, model, criterion)
logging.info('Valid_acc_top1: %f', valid_acc_top1)
logging.info('Valid_acc_top5: %f', valid_acc_top5)
epoch_duration = time.time() - epoch_start
logging.info('Epoch time: %ds.', epoch_duration)
is_best = False
if valid_acc_top5 > best_acc_top5:
best_acc_top5 = valid_acc_top5
if valid_acc_top1 > best_acc_top1:
best_acc_top1 = valid_acc_top1
is_best = True
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_acc_top1': best_acc_top1,
'optimizer': optimizer.state_dict(),
}, is_best, args.save)
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
if args.parallel: # multi gpu
num_gpus = torch.cuda.device_count()
logging.info('num of gpu devices = %d' % num_gpus)
else: # single gpu
torch.cuda.set_device(args.gpu)
logging.info('gpu device = %d' % args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
if args.parallel:
model = nn.DataParallel(model).cuda()
else:
model = model.cuda()
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
traindir = os.path.join(args.data, 'train')
validdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_data = dset.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.2),
transforms.ToTensor(),
normalize,
]))
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=4)
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=4)
if args.lr_scheduler == 'step':
# DARTS code
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
args.decay_period,
gamma=args.gamma)
elif args.lr_scheduler == 'cosine' or args.lr_scheduler == 'linear':
# PCDARTS code
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, float(args.epochs))
else:
raise ValueError("Wrong learning rate scheduler")
# ---- resume ---- #
start_epoch = 0
best_acc_top1 = 0.0
best_acc_top5 = 0.0
best_acc_epoch = 0
if args.resume:
# in multi-gpu???
if os.path.isfile(args.resume):
logging.info("=> loading checkpoint {}".format(args.resume))
device = torch.device("cuda")
checkpoint = torch.load(args.resume, map_location=device)
start_epoch = checkpoint['epoch']
best_acc_top1 = checkpoint['best_acc_top1']
best_acc_top5 = checkpoint['best_acc_top5']
best_acc_epoch = checkpoint['best_acc_epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logging.info(
"=> loaded checkpoint {} (trained until epoch {})".format(
args.resume, start_epoch - 1))
else:
raise ValueError("Wrong args.resume")
else:
logging.info("=> training from scratch")
for epoch in range(start_epoch, args.epochs):
scheduler.step()
if args.lr_scheduler == 'cosine' or args.lr_scheduler == 'step':
scheduler.step()
current_lr = scheduler.get_lr()[0]
elif args.lr_scheduler == 'linear':
current_lr = adjust_lr(optimizer, epoch)
if epoch < 5 and args.batch_size > 256:
for param_group in optimizer.param_groups:
param_group['lr'] = args.learning_rate * (epoch + 1) / 5.0
logging.info('Warming-up epoch: %d, LR: %e', epoch,
lr * (epoch + 1) / 5.0)
else:
logging.info('epoch %d lr %e', epoch, current_lr)
if args.parallel:
model.module.drop_path_prob = args.drop_path_prob * epoch / args.epochs
else:
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
epoch_start = time.time()
train_acc, train_obj = train(train_queue, model, criterion_smooth,
optimizer)
logging.info('train_acc %f', train_acc)
valid_acc_top1, valid_acc_top5, valid_obj = infer(
valid_queue, model, criterion)
is_best = (valid_acc_top1 > best_acc_top1)
if is_best:
best_acc_top1 = valid_acc_top1
best_acc_top5 = valid_acc_top5
best_acc_epoch = epoch + 1
utils.save(model, os.path.join(args.save, 'best_weights.pt'))
logging.info('valid_acc %f %f, best_acc %f %f (at epoch %d)',
valid_acc_top1, valid_acc_top5, best_acc_top1,
best_acc_top5, best_acc_epoch)
logging.info('epoch time %d sec.', time.time() - epoch_start)
utils.save_checkpoint(
{
'epoch': epoch + 1,
'best_acc_top1': best_acc_top1,
'best_acc_top5': best_acc_top5,
'best_acc_epoch': best_acc_epoch,
'state_dict': model.state_dict(),
'best_acc_top1': best_acc_top1,
'optimizer': optimizer.state_dict(),
}, is_best, args.save)
utils.save(model, os.path.join(args.save, 'weights.pt'))
# Horovod: (optional) compression algorithm.
compression = hvd.Compression.fp16 if args.fp16_allreduce else hvd.Compression.none
# Horovod: wrap optimizer with DistributedOptimizer.
optimizer = hvd.DistributedOptimizer(optimizer,
named_parameters=model.named_parameters(),
compression=compression)
# Restore from a previous checkpoint, if initial_epoch is specified.
# Horovod: restore on the first worker which will broadcast weights to other workers.
if resume_from_epoch > 0 and hvd.rank() == 0:
filepath = args.checkpoint_format.format(exp=args.save,
epoch=resume_from_epoch)
checkpoint = torch.load(filepath)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
# Horovod: broadcast parameters & optimizer state.
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)
def train(epoch, criterion):
model.train()
train_sampler.set_epoch(epoch)
train_loss = Metric('train_loss')
train_accuracy = Metric('train_accuracy')
if torch.cuda.current_device() == 0:
for param_group in optimizer.param_groups:
tmp_lr = param_group['lr']
def main():
if not torch.cuda.is_available():
logging.info('No GPU device available')
sys.exit(1)
num_gpus = torch.cuda.device_count()
args.gpu = args.local_rank % num_gpus
torch.cuda.set_device(args.gpu)
np.random.seed(args.seed)
cudnn.benchmark = True
cudnn.deterministic = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info("args = %s", args)
logging.info("unparsed_args = %s", unparsed)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.batch_size = args.batch_size // args.world_size
genotype = eval("genotypes.%s" % args.arch)
logging.info('---------Genotype---------')
logging.info(genotype)
logging.info('--------------------------')
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
model = model.cuda(args.gpu)
model = apex.parallel.DistributedDataParallel(model, delay_allreduce=True)
model_profile = Network(args.init_channels, CLASSES, args.layers,
args.auxiliary, genotype)
model_profile = model_profile.cuda(args.gpu)
model_input_size_imagenet = (1, 3, 224, 224)
model_profile.drop_path_prob = 0
flops, _ = profile(model_profile, model_input_size_imagenet)
logging.info("flops = %fMB, param size = %fMB", flops,
count_parameters_in_MB(model))
criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Prepare data
total_iters = per_epoch_iters * args.epochs
train_loader = get_train_dataloader(args.train_dir, args.batch_size,
args.local_rank, total_iters)
train_dataprovider = DataIterator(train_loader)
val_loader = get_val_dataloader(args.test_dir)
val_dataprovider = DataIterator(val_loader)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, float(args.epochs))
start_epoch = 0
best_acc_top1 = 0
best_acc_top5 = 0
checkpoint_tar = os.path.join(args.save, 'checkpoint.pth.tar')
if os.path.exists(checkpoint_tar):
logging.info('loading checkpoint {} ..........'.format(checkpoint_tar))
checkpoint = torch.load(
checkpoint_tar,
map_location={'cuda:0': 'cuda:{}'.format(args.local_rank)})
start_epoch = checkpoint['epoch'] + 1
model.load_state_dict(checkpoint['state_dict'])
logging.info("loaded checkpoint {} epoch = {}".format(
checkpoint_tar, checkpoint['epoch']))
# evaluation mode
if args.eval:
if args.eval_resume is not None:
checkpoint = torch.load(args.eval_resume)
model.module.drop_path_prob = 0
model.load_state_dict(checkpoint['state_dict'])
valid_acc_top1, valid_acc_top5 = infer(val_dataprovider,
model.module, val_iters)
print('valid_acc_top1: {}'.format(valid_acc_top1))
exit(0)
for epoch in range(start_epoch, args.epochs):
if args.lr_scheduler == 'cosine':
scheduler.step()
current_lr = scheduler.get_lr()[0]
elif args.lr_scheduler == 'linear':
current_lr = adjust_lr(optimizer, epoch)
else:
logging.info('Wrong lr type, exit')
sys.exit(1)
logging.info('Epoch: %d lr %e', epoch, current_lr)
if epoch < 5 and args.batch_size > 256:
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr * (epoch + 1) / 5.0
logging.info('Warming-up Epoch: %d, LR: %e', epoch,
current_lr * (epoch + 1) / 5.0)
model.module.drop_path_prob = args.drop_path_prob * epoch / args.epochs
epoch_start = time.time()
train_acc, train_obj = train(train_dataprovider, model,
criterion_smooth, optimizer,
per_epoch_iters)
writer.add_scalar('Train/Loss', train_obj, epoch)
writer.add_scalar('Train/LR', current_lr, epoch)
if args.local_rank == 0 and (epoch % 5 == 0
or args.epochs - epoch < 10):
valid_acc_top1, valid_acc_top5 = infer(val_dataprovider,
model.module, val_iters)
is_best = False
if valid_acc_top5 > best_acc_top5:
best_acc_top5 = valid_acc_top5
if valid_acc_top1 > best_acc_top1:
best_acc_top1 = valid_acc_top1
is_best = True
logging.info('Valid_acc_top1: %f', valid_acc_top1)
logging.info('Valid_acc_top5: %f', valid_acc_top5)
logging.info('best_acc_top1: %f', best_acc_top1)
epoch_duration = time.time() - epoch_start
logging.info('Epoch time: %ds.', epoch_duration)
save_checkpoint_(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'best_acc_top1': best_acc_top1,
'optimizer': optimizer.state_dict(),
}, args.save)
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
#torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled=True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device num = %d' % args.ngpu)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary, genotype, args.residual_wei, args.shrink_channel)
if args.parallel:
model = nn.DataParallel(model).cuda()
#model = nn.parallel.DistributedDataParallel(model).cuda()
else:
model = model.cuda()
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
optimizer = torch.optim.SGD(
#model.parameters(),
utils.set_group_weight(model, args.bn_no_wd, args.bias_no_wd),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay
)
resume = os.path.join(args.save, 'checkpoint.pth.tar')
if os.path.exists(resume):
print("=> loading checkpoint %s" % resume)
#checkpoint = torch.load(resume)
checkpoint = torch.load(resume, map_location = lambda storage, loc: storage.cuda(0))
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
#optimizer.load_state_dict(checkpoint['optimizer'])
optimizer.state_dict()['state'] = checkpoint['optimizer']['state']
print('=> loaded checkpoint epoch %d' % args.start_epoch)
if args.start_epoch >= args.epochs:
print('training finished')
sys.exit(0)
traindir = os.path.join(args.data, 'train')
validdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
train_data = dset.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(args.image_size),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(
brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.1),
transforms.ToTensor(),
normalize,
]))
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(int((256.0 / 224) * args.image_size)),
transforms.CenterCrop(args.image_size),
transforms.ToTensor(),
normalize,
]))
train_queue = torch.utils.data.DataLoader(
train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=nworker)
valid_queue = torch.utils.data.DataLoader(
valid_data, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=nworker)
best_acc_top1 = 0
for epoch in range(args.start_epoch, args.epochs):
if args.lr_strategy == 'cos':
lr = utils.set_lr(optimizer, epoch, args.epochs, args.learning_rate)
#elif args.lr_strategy == 'step':
# scheduler.step()
# lr = scheduler.get_lr()[0]
logging.info('epoch %d lr %e', epoch, lr)
if args.parallel:
model.module.drop_path_prob = args.drop_path_prob * epoch / args.epochs
else:
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
train_acc, train_obj = train(train_queue, model, criterion_smooth, optimizer, epoch)
logging.info('train_acc %f', train_acc)
utils.save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_acc_top1': train_acc,
'optimizer' : optimizer.state_dict(),
}, False, args.save)
#if epoch >= args.early_stop:
# break
valid_acc_top1, valid_acc_top5, valid_obj = infer(valid_queue, model, criterion)
logging.info('valid_acc_top1 %f', valid_acc_top1)
logging.info('valid_acc_top5 %f', valid_acc_top5)
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
cudnn.benchmark = True
cudnn.enabled = True
logging.info("args = %s", args)
logging.info('Training with config:')
logging.info(pprint.pformat(config))
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
model.drop_path_prob = args.drop_path_prob
model = nn.DataParallel(model)
model = model.cuda()
model.load_state_dict(torch.load(args.model_path))
imagenet = imagenet_data.ImageNet12(
trainFolder=os.path.join(args.data_path, 'train'),
testFolder=os.path.join(args.data_path, 'val'),
num_workers=config.data.num_workers,
data_config=config.data)
valid_queue = imagenet.getTestLoader(config.data.batch_size)
trainer = Trainer(None, valid_queue, None, config, args.report_freq)
with torch.no_grad():
val_acc_top1, val_acc_top5, valid_obj, batch_time = trainer.infer(
model)
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
print(torch.cuda.device_count())
np.random.seed(args.seed)
cudnn.benchmark = True
torch.manual_seed(args.seed)
if hvd.rank() == 0:
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
best_acc_top1 = 0
start_epoch = 0
if args.resume:
checkpoint = torch.load(os.path.join(args.save, 'checkpoint.pth.tar'))
best_checkpoint = torch.load(
os.path.join(args.save, 'model_best.pth.tar'))
start_epoch = checkpoint['epoch']
best_acc_top1 = best_checkpoint['best_acc_top1']
start_epoch = hvd.broadcast(torch.tensor(start_epoch),
root_rank=0,
name='start_epoch').item()
best_acc_top1 = hvd.broadcast(torch.tensor(best_acc_top1),
root_rank=0,
name='best_acc_top1').item()
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
if args.parallel:
model = nn.DataParallel(model).cuda()
else:
model = model.cuda()
if hvd.rank() == 0:
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate * hvd.size(),
momentum=args.momentum,
weight_decay=args.weight_decay)
# ***************** horovod *******************
optimizer = hvd.DistributedOptimizer(
optimizer, named_parameters=model.named_parameters())
# ***************** horovod *******************
traindir = os.path.join(args.data, 'train')
validdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_data = dset.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.2),
transforms.ToTensor(),
normalize,
]))
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_data, num_replicas=hvd.size(), rank=hvd.rank())
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
pin_memory=True,
num_workers=args.num_workers,
sampler=train_sampler)
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=args.batch_size,
shuffle=False,
pin_memory=True,
num_workers=args.num_workers)
if start_epoch > 0 and hvd.rank() == 0:
checkpoint = torch.load(os.path.join(args.save, 'checkpoint.pth.tar'))
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("checkpoint {}, model, optimizer was loaded".format(start_epoch))
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)
if not args.resume:
set_lr(0, 0, len(train_queue), optimizer, args.scheduler)
for epoch in range(start_epoch, args.epochs + args.warmup_epochs):
if hvd.rank() == 0:
lr = optimizer.param_groups[0]['lr']
logging.info('epoch %d lr %e', epoch, lr)
with open(os.path.join(args.save, 'learning_rate.txt'),
mode='a') as f:
f.write(str(lr) + '\n')
if args.parallel:
model.module.drop_path_prob = args.drop_path_prob * epoch / args.epochs
else:
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
train_acc, train_obj = train(train_queue, train_sampler, model,
criterion_smooth, optimizer, epoch)
if hvd.rank() == 0:
logging.info('train_acc %f', train_acc)
with open(os.path.join(args.save, "train_acc.txt"), mode='a') as f:
f.write(str(train_acc) + '\n')
with open(os.path.join(args.save, "train_loss.txt"),
mode='a') as f:
f.write(str(train_obj) + '\n')
valid_acc_top1, valid_acc_top5, valid_obj = infer(
valid_queue, model, criterion)
if hvd.rank() == 0:
logging.info('valid_acc_top1 %f', valid_acc_top1)
logging.info('valid_acc_top5 %f', valid_acc_top5)
with open(os.path.join(args.save, "test_acc_1.txt"),
mode='a') as f:
f.write(str(valid_acc_top1) + '\n')
with open(os.path.join(args.save, "test_acc_5.txt"),
mode='a') as f:
f.write(str(valid_acc_top5) + '\n')
with open(os.path.join(args.save, "test_loss.txt"), mode='a') as f:
f.write(str(valid_obj) + '\n')
is_best = False
if valid_acc_top1 > best_acc_top1:
best_acc_top1 = valid_acc_top1
is_best = True
if hvd.rank() == 0:
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_acc_top1': best_acc_top1,
'optimizer': optimizer.state_dict(),
}, is_best, args.save)