code: def main():  # noqa: C901
      world_size = 1
  
      args = parse_args()
      B = args.batch_size
      H, W = args.height, args.width
  
      img = torch.randn(args.steps * B, 3, H, W)
      labels = torch.arange(B).repeat(args.steps)
      print(img.sum())
  
      train_dataset = TensorDataset(img, labels)
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=B,
          num_workers=args.workers,
          pin_memory=True,
      )
  
      if not args.disable_dp:
          model = models.__dict__[args.architecture](
              pretrained=False,
              norm_layer=(lambda c: nn.GroupNorm(args.gn_groups, c)),
          )
      else:
          model = models.__dict__[args.architecture](pretrained=False)
  
      model = model.to(args.device)
      print(
          "Model size: "
          + pretty_number(sum([p.numel() for p in model.parameters()]))
      )
  
      # Use the right distributed module wrapper if distributed training is enabled
      if world_size > 1:
          if not args.disable_dp:
              if args.dist_algo == "naive":
                  model = DPDDP(model)
              elif args.dist_algo == "ddp_hook":
                  model = DDP(model, device_ids=[args.device])
              else:
                  raise NotImplementedError(
                      f"Unrecognized argument for the distributed algorithm: {args.dist_algo}"
                  )
          else:
              model = DDP(model, device_ids=[args.device])
  
      if args.optim == "SGD":
          optimizer = optim.SGD(
              model.parameters(),
              lr=args.lr,
              momentum=args.momentum,
              weight_decay=args.weight_decay,
          )
      elif args.optim == "RMSprop":
          optimizer = optim.RMSprop(model.parameters(), lr=args.lr)
      elif args.optim == "Adam":
          optimizer = optim.Adam(model.parameters(), lr=args.lr)
      else:
          raise NotImplementedError(
              "Optimizer not recognized. Please check spelling"
          )
  
      privacy_engine = None
      if not args.disable_dp:
          if args.clip_per_layer:
              # Each layer has the same clipping threshold. The total grad norm is still bounded by `args.max_per_sample_grad_norm`.
              n_layers = len(
                  [
                      (n, p)
                      for n, p in model.named_parameters()
                      if p.requires_grad
                  ]
              )
              max_grad_norm = [
                  args.max_per_sample_grad_norm / np.sqrt(n_layers)
              ] * n_layers
          else:
              max_grad_norm = args.max_per_sample_grad_norm
  
          privacy_engine = PrivacyEngine(
              secure_mode=args.secure_mode,
          )
          model, optimizer, train_loader = privacy_engine.make_private(
              module=model,
              optimizer=optimizer,
              data_loader=train_loader,
              noise_multiplier=args.sigma,
              max_grad_norm=max_grad_norm,
              poisson_sampling=False,
          )
  
      criterion = nn.CrossEntropyLoss()
      model.train()
      print(type(model))
  
      if args.benchmark_data_loader:
          torch.cuda.synchronize()
          start = time.time()
          data_time = 0
          data_end = time.time()
          for images, target in tqdm(train_loader):
              data_time += time.time() - data_end
              images = images.cuda(non_blocking=True)
  
              # compute output
              output = model(images)
              loss = criterion(output, target.cuda(non_blocking=True))
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
              data_end = time.time()
      else:
          images = torch.randn(B, 3, H, W).cuda()
          target = torch.arange(B).cuda()
  
          torch.cuda.synchronize()
          start = time.time()
          for _ in tqdm(range(args.steps)):
              output = model(images)
              loss = criterion(output, target)
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
  
      torch.cuda.synchronize()
      elapsed = time.time() - start
      if args.benchmark_data_loader:
          elapsed -= data_time
          print(f"Data time {data_time:.2f}")
  
      print(f"Took {elapsed:.2f}")
      speed = args.steps * args.batch_size / elapsed
      print(f"Speed: {speed:.2f} img/s")
  
file_path: opacus/examples/vision_benchmark.py
imports: import argparse
  import time
  import numpy
  import torch
  import torch.nn
  import torch.optim
  import torch.utils.data
  import torch.utils.data.distributed
  from opacus import PrivacyEngine
  from opacus.distributed import DifferentiallyPrivateDistributedDataParallel
  from torch.nn.parallel import DistributedDataParallel
  from torch.utils.data import TensorDataset
  from torchvision import models
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types