code: def main(args):
      quant_modules.initialize()
      assert torch.cuda.is_available(), "only support GPU!"
  
      (
          train_images_path,
          train_images_label,
          val_images_path,
          val_images_label,
      ) = read_split_data(args.data_path)
  
      data_transform = {
          "train": transforms.Compose(
              [
                  transforms.RandomResizedCrop(224),
                  transforms.RandomHorizontalFlip(),
                  transforms.ToTensor(),
                  transforms.Normalize(
                      [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
                  ),
              ]
          ),
          "val": transforms.Compose(
              [
                  transforms.Resize(256),
                  transforms.CenterCrop(224),
                  transforms.ToTensor(),
                  transforms.Normalize(
                      [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
                  ),
              ]
          ),
      }
  
      # 实例化训练数据集
      train_dataset = MyDataSet(
          images_path=train_images_path,
          images_class=train_images_label,
          transform=data_transform["train"],
      )
  
      # 实例化验证数据集
      val_dataset = MyDataSet(
          images_path=val_images_path,
          images_class=val_images_label,
          transform=data_transform["val"],
      )
  
      batch_size = args.batch_size
      nw = min(
          [os.cpu_count(), batch_size if batch_size > 1 else 0, 8]
      )  # number of workers
      print("Using {} dataloader workers every process".format(nw))
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=batch_size,
          shuffle=True,
          pin_memory=True,
          num_workers=nw,
          collate_fn=train_dataset.collate_fn,
      )
  
      val_loader = torch.utils.data.DataLoader(
          val_dataset,
          batch_size=batch_size,
          shuffle=False,
          pin_memory=True,
          num_workers=nw,
          collate_fn=val_dataset.collate_fn,
      )
  
      # ########################## #
      # Post Training Quantization #
      # ########################## #
      # We will use histogram based calibration for activations and the default max calibration for weights.
      quant_desc_input = QuantDescriptor(calib_method="histogram")
      quant_nn.QuantConv2d.set_default_quant_desc_input(
          quant_desc_input
      )
      quant_nn.QuantLinear.set_default_quant_desc_input(
          quant_desc_input
      )
  
      model = create_model(num_classes=args.num_classes)
      assert os.path.exists(
          args.weights
      ), "weights file: '{}' not exist.".format(args.weights)
      model.load_state_dict(
          torch.load(args.weights, map_location="cpu")
      )
      model.cuda()
  
      # It is a bit slow since we collect histograms on CPU
      with torch.no_grad():
          collect_stats(
              model, val_loader, num_batches=1000 // batch_size
          )
          compute_amax(model, method="percentile", percentile=99.99)
          # validate
          evaluate(model=model, data_loader=val_loader, epoch=0)
  
      torch.save(model.state_dict(), "quant_model_calibrated.pth")
  
      if args.qat:
          # ########################### #
          # Quantization Aware Training #
          # ########################### #
          pg = [p for p in model.parameters() if p.requires_grad]
          optimizer = optim.SGD(
              pg, lr=args.lr, momentum=0.9, weight_decay=5e-5
          )
          # Scheduler(half of a cosine period)
          lf = (
              lambda x: (math.cos(x * math.pi / 2 / args.epochs))
              * (1 - args.lrf)
              + args.lrf
          )
          scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
  
          for epoch in range(args.epochs):
              # train
              train_one_epoch(
                  model=model,
                  optimizer=optimizer,
                  data_loader=train_loader,
                  epoch=epoch,
              )
  
              scheduler.step()
  
              # validate
              evaluate(model=model, data_loader=val_loader, epoch=epoch)
  
      export_onnx(model, args.onnx_filename, args.onnx_bs)
  
file_path: deep-learning-for-image-processing/deploying_service/deploying_pytorch/convert_tensorrt/convert_resnet34/quantization.py
imports: import os
  import math
  import argparse
  from absl import logging
  from tqdm import tqdm
  import torch
  import torch.optim
  import torch.optim.lr_scheduler
  from torchvision import transforms
  from torchvision.models.resnet import resnet34
  from pytorch_quantization import nn
  from pytorch_quantization import quant_modules, calib
  from pytorch_quantization.tensor_quant import QuantDescriptor
  from my_dataset import MyDataSet
  from utils import read_split_data, train_one_epoch, evaluate
label_desc: Use image node_id 4 for calling a global function with example usage: main(args) without return types