docstr: Test the behavior of max_depth on a layered model like ResNet
docstr_tok: ["Test","the","behavior","of","max_depth","on","a","layered","model","like","ResNet"]
code: def test_resnet_max_depth(self) -> None:
      """Test the behavior of max_depth on a layered model like ResNet"""
      pretrained_model = models.resnet.resnet18(
          weights=ResNet18_Weights.IMAGENET1K_V1
      )
  
      # max_depth = None
      ms1 = get_module_summary(pretrained_model)
  
      self.assertEqual(len(ms1.submodule_summaries), 10)
      self.assertEqual(
          len(ms1.submodule_summaries["layer2"].submodule_summaries), 2
      )
      self.assertEqual(
          len(
              ms1.submodule_summaries["layer2"]
              .submodule_summaries["layer2.0"]
              .submodule_summaries
          ),
          6,
      )
  
      ms2 = get_summary_and_prune(pretrained_model, max_depth=1)
      self.assertEqual(len(ms2.submodule_summaries), 0)
      self.assertNotIn("layer2", ms2.submodule_summaries)
  
      ms3 = get_summary_and_prune(pretrained_model, max_depth=2)
      self.assertEqual(len(ms3.submodule_summaries), 10)
      self.assertEqual(
          len(ms1.submodule_summaries["layer2"].submodule_summaries), 2
      )
      self.assertNotIn(
          "layer2.0",
          ms3.submodule_summaries["layer2"].submodule_summaries,
      )
      inp = torch.randn(1, 3, 224, 224)
      ms4 = get_summary_and_prune(
          pretrained_model, max_depth=2, module_args=(inp,)
      )
  
      self.assertEqual(len(ms4.submodule_summaries), 10)
      self.assertEqual(ms4.flops_forward, 1814073344)
      self.assertEqual(ms4.flops_backward, 3510132736)
      self.assertEqual(
          ms4.submodule_summaries["layer2"].flops_forward, 411041792
      )
      self.assertEqual(
          ms4.submodule_summaries["layer2"].flops_backward, 822083584
      )
  
      # These should stay constant for all max_depth values
      ms_list = [ms1, ms2, ms3, ms4]
      for ms in ms_list:
          self.assertEqual(ms.module_name, "")
          self.assertEqual(ms.module_type, "ResNet")
          self.assertEqual(ms.num_parameters, 11689512)
          self.assertEqual(ms.num_trainable_parameters, 11689512)
          self.assertFalse(ms.has_uninitialized_param)
  
imports: import unittest
  from typing import Any, Dict, Optional, Tuple
  import torch
  from torchtnt.utils.module_summary import _get_human_readable_count, get_module_summary, ModuleSummary, prune_module_summary
  from torchvision import models
  from torchvision.models.alexnet import AlexNet_Weights
  from torchvision.models.resnet import ResNet18_Weights
label_desc: Use image node_id 7 for calling the ModuleSummaryTest obj's underlying member method code with example usage: obj.test_resnet_max_depth() without return types