text: import time
  
  import numpy as np
  import torch
  import torch.nn.functional as F
  import wandb
  from torch import nn
  from torch.nn import CrossEntropyLoss
  from torch.nn.utils import prune
  from torchvision.models import inception, inception_v3, SqueezeNet
  from tqdm import tqdm
  
  import os
  import simplify
  from tests.benchmark_models import models
  
  device = os.environ.get('DEVICE', 'cuda')
  
  
  def time_model(model, x, y):
      forward_time = []
      backward_time = []
      model.to(device)
      for j in range(10):  # tqdm(range(10), desc="Pruned test"):
          if device == "cuda":
              starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
              starter.record()
          else:
              start = time.perf_counter()
          
          output = model(x)  # FORWARD PASS
          
          if device == "cuda":
              ender.record()
              torch.cuda.synchronize()
          else:
              end = time.perf_counter()
          
          forward_time.append(starter.elapsed_time(ender) if device == ("cuda") else end - start)
          
          loss = F.cross_entropy(output, y)
          
          if device == "cuda":
              starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
              starter.record()
          else:
              start = time.perf_counter()
          
          loss.backward()  # BACKWARD PASS
          
          if device == "cuda":
              ender.record()
              torch.cuda.synchronize()
          else:
              end = time.perf_counter()
          
          backward_time.append(starter.elapsed_time(ender) if device == "cuda" else end - start)
      
      forward_time = forward_time[1:]
      backward_time = backward_time[1:]
      
      return forward_time, backward_time
  
  
  def main(network):
      print('=> Benchmarking', network.__name__)
      if network.__name__ in ["inception_v3", "googlenet"]:
          model = network(False, aux_logits=False)
      else:
          model = network(False)
      
      batch_size = 32
      h, w = 224, 224
      if network.__name__ == "inception_v3":
          h, w = 299, 299
          
      fake_input = torch.randint(0, 256, (batch_size, 3, h, w))
      fake_input = fake_input.float() / 255.
      fake_target = torch.randint(0, 1000, (batch_size,)).long()
      
      fake_input, fake_target = fake_input.to(device), fake_target.to(device)
      
      criterion = CrossEntropyLoss()
      
      prune_step = 0.05
      iterations = int(1. / prune_step)
      
      total_neurons = 0
      remaining_neurons = 0
      for name, module in model.named_modules():
          if isinstance(module, nn.Conv2d):
              total_neurons += module.weight.shape[0]
              remaining_neurons += module.weight.shape[0]
      
      x = []
      pruned_y_forward = []
      pruned_y_forward_std = []
      pruned_y_backward = []
      pruned_y_backward_std = []
      
      simplified_y_forward = []
      simplified_y_forward_std = []
      simplified_y_backward = []
      simplified_y_backward_std = []
      
      amount = 0.
      
      wandb.init(
          config={'arch': network.__name__, 'device': os.environ.get('DEVICE', 'cuda')},
          group="benchmark_passes"
      )
      
      for i in tqdm(range(iterations), desc="Benchmark"):
          if amount > 1.:
              break
          if network.__name__ in ["inception_v3", "googlenet"]:
              model = network(False, aux_logits=False)
          else:
              model = network(False)
          
          _, _ = time_model(model, fake_input, fake_target)
  
          # First loop is the full model
          if i > 0:
              remaining_neurons = 0
              for name, module in model.named_modules():
                  if isinstance(model, SqueezeNet) and 'classifier.1' in name:
                      continue
                  if isinstance(module, nn.Conv2d):
                      prune.ln_structured(module, 'weight', amount=amount, n=2, dim=0)
                      ch_sum = module.weight.sum(dim=(1, 2, 3))
                      remaining_neurons += ch_sum[ch_sum != 0].shape[0]
                      prune.remove(module, 'weight')
          
          x.append(100 - (remaining_neurons / total_neurons) * 100)
          
          # PRUNED
          forward_time, backward_time = time_model(model, fake_input, fake_target)
          
          pruned_y_forward.append(np.mean(forward_time))
          pruned_y_forward_std.append(np.std(forward_time))
          pruned_y_backward.append(np.mean(backward_time))
          pruned_y_backward_std.append(np.std(backward_time))
          
          wandb.log({
              'pruned.forward':      np.mean(forward_time),
              'pruned.forward_std':  np.std(forward_time),
              'pruned.backward':     np.mean(backward_time),
              'pruned.backward_std': np.std(backward_time)
          }, commit=False)
          
          # SIMPLIFIED
          model.eval()
          pinned_out = None
          if 'densenet' in network.__name__ or 'squeezenet' in network.__name__:
              pinned_out = [] 
          simplify.simplify(model, torch.zeros(1, 3, h, w).to(device), 
                            fuse_bn=False, training=True, pinned_out=pinned_out)
          model.train()
          
          forward_time, backward_time = time_model(model, fake_input, fake_target)
          
          simplified_y_forward.append(np.mean(forward_time))
          simplified_y_forward_std.append(np.std(forward_time))
          simplified_y_backward.append(np.mean(backward_time))
          simplified_y_backward_std.append(np.std(backward_time))
          
          wandb.log({
              'simplified.forward':      np.mean(forward_time),
              'simplified.forward_std':  np.std(forward_time),
              'simplified.backward':     np.mean(backward_time),
              'simplified.backward_std': np.std(backward_time),
              'remaining_neurons':       100. - x[-1],
              'amount': amount
          })
          
          amount += prune_step
          del model
          torch.cuda.empty_cache()
      
      wandb.finish()
  
  
  if __name__ == '__main__':
      for arch in models:
          main(arch)
  