code: def main(config):
      # Training settings
      parser = argparse.ArgumentParser(
          description="PyTorch MNIST Example"
      )
      parser.add_argument(
          "--batch-size",
          type=int,
          default=64,
          metavar="N",
          help="input batch size for training (default: 64)",
      )
      parser.add_argument(
          "--test-batch-size",
          type=int,
          default=1000,
          metavar="N",
          help="input batch size for testing (default: 1000)",
      )
      parser.add_argument(
          "--epochs",
          type=int,
          default=1,
          metavar="N",
          help="number of epochs to train (default: 14)",
      )
      parser.add_argument(
          "--lr",
          type=float,
          default=1.0,
          metavar="LR",
          help="learning rate (default: 1.0)",
      )
      parser.add_argument(
          "--gamma",
          type=float,
          default=0.7,
          metavar="M",
          help="Learning rate step gamma (default: 0.7)",
      )
      parser.add_argument(
          "--no-cuda",
          action="store_true",
          default=False,
          help="disables CUDA training",
      )
      parser.add_argument(
          "--dry-run",
          action="store_true",
          default=False,
          help="quickly check a single pass",
      )
      parser.add_argument(
          "--seed",
          type=int,
          default=1,
          metavar="S",
          help="random seed (default: 1)",
      )
      parser.add_argument(
          "--log-interval",
          type=int,
          default=10,
          metavar="N",
          help="how many batches to wait before logging training status",
      )
      args, _ = parser.parse_known_args()
  
      use_cuda = not args.no_cuda and torch.cuda.is_available()
  
      torch.manual_seed(args.seed)
  
      device = torch.device("cuda" if use_cuda else "cpu")
  
      kwargs = {"batch_size": args.batch_size}
      if use_cuda:
          kwargs.update(
              {"num_workers": 1, "pin_memory": True, "shuffle": True},
          )
  
      def new_getitem(self, index):
          img, target = self.data[index], int(self.targets[index])
  
          # doing this so that it is consistent with all other datasets
          # to return a PIL Image
          img = Image.fromarray(img.numpy(), mode="L")
          img = img.convert("RGB")
  
          if self.transform is not None:
              img = self.transform(img)
  
          if self.target_transform is not None:
              target = self.target_transform(target)
  
          return img, target
  
      datasets.mnist.MNIST.__getitem__ = new_getitem
  
      transform = transforms.Compose(
          [
              transforms.Resize(256),
              transforms.CenterCrop(224),
              transforms.ToTensor(),
              transforms.Normalize(
                  mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
              ),
          ]
      )
      dataset1 = datasets.MNIST(
          "/tmp/torch_mnist_data",
          train=True,
          download=True,
          transform=transform,
      )
      dataset2 = datasets.MNIST(
          "/tmp/torch_mnist_data", train=False, transform=transform
      )
      train_loader = torch.utils.data.DataLoader(dataset1, **kwargs)
      test_loader = torch.utils.data.DataLoader(dataset2, **kwargs)
  
      model = torch.hub.load(
          "pytorch/vision:v{}".format(torchvision.__version__),
          "mobilenet_v2",
          pretrained=True,
      ).to(device)
      optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
      scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
      _, test_acc = test(model, device, test_loader)
  
      def evaluator(val_loader, model):
          _, test_acc = test(model, device, val_loader)
          return test_acc
  
      compressor = aup.compression.create_compressor(
          model,
          config,
          evaluator=evaluator,
          val_loader=test_loader,
          flops_ratio=0.5,
          model_type="mobilenetv2",
          n_calibration_batches=1,
          n_points_per_layer=1,
          warmup=1,
          hidden1=32,
          hidden2=32,
          train_episode=1,
          max_episode_length=1,
      )
      compressor.compress()
  
      flops, _, _ = compressor.count_flops_params((1, 3, 28, 28))
      print("test_acc={} flops={}".format(test_acc, flops))
  
      if "save_model" in config and config["save_model"]:
          compressor.export_model(
              model_path="mnist_compressed.pth",
              mask_path="mnist_compressed_mask.pth",
              save_best_only=True,
              folder_name=config["folder_name"],
          )
  
      aup.print_result((test_acc - 1.0) / np.log(flops))
  
file_path: auptimizer/Examples/compression/mnist_pytorch/mnist_amc.py
imports: from __future__ import print_function
  import argparse
  import sys
  import time
  import numpy
  from PIL import Image
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import torchvision
  from torchvision import datasets, transforms
  from torch.optim.lr_scheduler import StepLR
  import aup
label_desc: Use image node_id 3 for calling a global function with example usage: main(config) and returns: img,  target, test_acc