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)