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 write_pytorch_weight(self):
model = NetworkImageNet(216, 1001, 12, False, PNASNet)
model.drop_path_prob = 0
model.eval()
self.used_keys = []
self.convert_conv(model.conv0, 'conv0/weights')
self.convert_bn(model.conv0_bn, 'conv0_bn/gamma', 'conv0_bn/beta',
'conv0_bn/moving_mean', 'conv0_bn/moving_variance')
self.convert_cell(model.stem1, 'cell_stem_0/')
self.convert_cell(model.stem2, 'cell_stem_1/')
for i in range(12):
self.convert_cell(model.cells[i], 'cell_{}/'.format(i))
self.convert_fc(model.classifier, 'final_layer/FC/weights',
'final_layer/FC/biases')
print('Conversion complete!')
print('Check 1: whether all TF variables are used...')
assert len(self.weight_dict) == len(self.used_keys)
print('Pass!')
model = model.cuda()
image = self.tf_image_proc.transpose((2, 0, 1))
image = Variable(self.Tensor(image)).cuda()
logits, _ = model(image.unsqueeze(0))
self.pytorch_logits = logits.data.cpu().numpy()
print('Check 2: whether logits have small diff...')
assert np.max(np.abs(self.tf_logits - self.pytorch_logits)) < 1e-5
print('Pass!')
model_path = 'data/PNASNet-5_Large.pth'
torch.save(model.state_dict(), model_path)
print('PyTorch model saved to {}'.format(model_path))
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("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, NUM_CLASSES, args.layers,
config.optim.auxiliary, genotype)
start_epoch = 0
model.eval()
model.drop_path_prob = args.drop_path_prob * 0
# compute the params as well as the multi-adds
params = count_parameters_in_MB(model)
logging.info("Params = %.2fMB" % params)
mult_adds = comp_multadds(model, input_size=config.data.input_size)
logging.info("Mult-Adds = %.2fMB" % mult_adds)
model.train()
if len(args.gpus) > 1:
model = nn.DataParallel(model)
model = model.cuda()
if config.optim.label_smooth:
criterion = CrossEntropyLabelSmooth(NUM_CLASSES,
config.optim.smooth_alpha)
else:
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
optimizer = torch.optim.SGD(model.parameters(),
config.optim.init_lr,
momentum=config.optim.momentum,
weight_decay=config.optim.weight_decay)
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,
type_of_data_augmentation=config.data.type_of_data_aug,
data_config=config.data,
size_images=config.data.input_size[1],
scaled_size=config.data.scaled_size[1])
train_queue, valid_queue = imagenet.getTrainTestLoader(
config.data.batch_size)
if config.optim.lr_schedule == 'cosine':
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, float(config.train_params.epochs))
trainer = Trainer(train_queue, valid_queue, criterion, config,
args.report_freq)
best_epoch = [0, 0, 0] # [epoch, acc_top1, acc_top5]
lr = config.optim.init_lr
for epoch in range(start_epoch, config.train_params.epochs):
if config.optim.lr_schedule == 'cosine':
scheduler.step()
current_lr = scheduler.get_lr()[0]
elif config.optim.lr_schedule == 'linear': # with warmup initial
optimizer, current_lr = adjust_lr(optimizer,
config.train_params.epochs, lr,
epoch)
else:
print('Wrong lr type, exit')
sys.exit(1)
if epoch < 5: # Warmup epochs for 5
current_lr = lr * (epoch + 1) / 5.0
for param_group in optimizer.param_groups:
param_group['lr'] = current_lr
logging.info('Warming-up Epoch: %d, LR: %e', epoch,
lr * (epoch + 1) / 5.0)
logging.info('Epoch: %d lr %e', epoch, current_lr)
if len(args.gpus) > 1:
model.module.drop_path_prob = args.drop_path_prob * epoch / config.train_params.epochs
else:
model.drop_path_prob = args.drop_path_prob * epoch / config.train_params.epochs
train_acc_top1, train_acc_top5, train_obj, batch_time, data_time = trainer.train(
model, optimizer, epoch)
with torch.no_grad():
val_acc_top1, val_acc_top5, batch_time, data_time = trainer.infer(
model, epoch)
if val_acc_top1 > best_epoch[1]:
best_epoch = [epoch, val_acc_top1, val_acc_top5]
if epoch >= 0: # 120
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.module.state_dict(),
'best_acc_top1': val_acc_top1,
'optimizer': optimizer.state_dict(),
},
save_path=args.save,
epoch=epoch,
is_best=True)
if len(args.gpus) > 1:
utils.save(
model.module.state_dict(),
os.path.join(
args.save,
'weights_{}_{}.pt'.format(epoch, val_acc_top1)))
else:
utils.save(
model.state_dict(),
os.path.join(
args.save,
'weights_{}_{}.pt'.format(epoch, val_acc_top1)))
logging.info('BEST EPOCH %d val_top1 %.2f val_top5 %.2f',
best_epoch[0], best_epoch[1], best_epoch[2])
logging.info(
'epoch: {} \t train_acc_top1: {:.4f} \t train_loss: {:.4f} \t val_acc_top1: {:.4f}'
.format(epoch, train_acc_top1, train_obj, val_acc_top1))
logging.info("Params = %.2fMB" % params)
logging.info("Mult-Adds = %.2fMB" % mult_adds)
