code: def test_memory_tracking(self) -> None:
      # Create a model with a hierarchy of modules
      torch.manual_seed(0)
      model = nn.Sequential(
          nn.Sequential(
              nn.Conv2d(3, 64, kernel_size=3),
              nn.BatchNorm2d(64),
              nn.ReLU(inplace=False),
              nn.AdaptiveAvgPool2d(output_size=(1, 1)),
          ),
          nn.Flatten(start_dim=1),
          nn.Sequential(nn.Linear(64, 2), nn.ReLU(inplace=True)),
      ).cuda()
  
      # Track a fake forward / backward
      tracker = LayerwiseMemoryTracker()
      tracker.monitor(model)
      x = torch.randn(size=(2, 3, 224, 224)).cuda()
      target = torch.LongTensor([0, 1]).cuda()
      criterion = nn.CrossEntropyLoss()
      criterion(model(x), target).backward()
  
      # Verify that only leaf modules are tracked
      tracked_names = {
          trace.module_name for trace in tracker.memory_traces
      }
      expected_names = {"0.0", "0.1", "0.2", "0.3", "1", "2.0", "2.1"}
      self.assertEqual(expected_names, tracked_names)
  
      # Verify that memory tracking for ReLU is sound
      self.assertEqual(
          25233408,
          tracker.forward_traces[2].event.memory_activations,
          "ReLU(inplace=False) should allocate activations",
      )
      self.assertEqual(
          0,
          tracker.forward_traces[6].event.memory_activations,
          "ReLU(inplace=True) should NOT allocate activations",
      )
  
      # Verify that overall memory tracking is sound
      summary = tracker.summary
      self.assertGreaterEqual(
          summary.total_forward_allocations,
          summary.total_activation_allocations,
      )
  
      top_act_producers = summary.top_forward_activation_producers[:3]
      self.assertEqual("0.0", top_act_producers[0].module_name)
      self.assertEqual("0.1", top_act_producers[1].module_name)
      self.assertEqual("0.2", top_act_producers[2].module_name)
      self.assertEqual(7168, top_act_producers[0].module_params)
      self.assertEqual(512, top_act_producers[1].module_params)
      self.assertEqual(0, top_act_producers[2].module_params)
      for trace in top_act_producers:
          self.assertEqual(25233408, trace.event.memory_activations)
  
file_path: vissl/tests/test_layer_memory_tracking.py
imports: import unittest
  import torch
  import torch.distributed
  import torch.multiprocessing
  import torch.nn
  import torchvision.models
  from fairscale.nn import FullyShardedDataParallel
  from vissl.utils.layer_memory_tracking import find_best_reset_points, LayerwiseMemoryTracker, ProcessGroupTracker
  from vissl.utils.test_utils import gpu_test, init_distributed_on_file, with_temp_files, with_timing
label_desc: Use image node_id 1 for calling the TestLayerMemoryTracking obj's underlying member method code with example usage: obj.test_memory_tracking() without return types