TylerYep / torchinfo

View model summaries in PyTorch!

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Running `torchinfo.summary` on YOLO from `ultralytics` package?

ego-thales opened this issue · comments

Hello,

I managed without problem to run torchinfo.summary on YOLOv5n that I loaded from the torch hub.

>>> torchinfo.summary(torch.hub.load('ultralytics/yolov5', 'yolov5n'), input_size=(1, 3, 640, 640), depth=4)
Using cache found in [...]/.cache/torch/hub/ultralytics_yolov5_master
YOLOv5 🚀 2024-5-28 Python-3.12.3 torch-2.3.0+cu121 CPU

Fusing layers... 
YOLOv5n summary: 213 layers, 1867405 parameters, 0 gradients, 4.5 GFLOPs
Adding AutoShape... 
=========================================================================================================
Layer (type:depth-idx)                                  Output Shape              Param #
=========================================================================================================
AutoShape                                               [1, 25200, 85]            --
├─DetectMultiBackend: 1-1                               [1, 25200, 85]            --
│    └─DetectionModel: 2-1                              [1, 25200, 85]            --
│    │    └─Sequential: 3-1                             --                        --
│    │    │    └─Conv: 4-1                              [1, 16, 320, 320]         (1,744)
│    │    │    └─Conv: 4-2                              [1, 32, 160, 160]         (4,640)
│    │    │    └─C3: 4-3                                [1, 32, 160, 160]         (4,704)
│    │    │    └─Conv: 4-4                              [1, 64, 80, 80]           (18,496)
│    │    │    └─C3: 4-5                                [1, 64, 80, 80]           (28,928)
│    │    │    └─Conv: 4-6                              [1, 128, 40, 40]          (73,856)
│    │    │    └─C3: 4-7                                [1, 128, 40, 40]          (156,288)
│    │    │    └─Conv: 4-8                              [1, 256, 20, 20]          (295,168)
│    │    │    └─C3: 4-9                                [1, 256, 20, 20]          (295,680)
│    │    │    └─SPPF: 4-10                             [1, 256, 20, 20]          (164,224)
│    │    │    └─Conv: 4-11                             [1, 128, 20, 20]          (32,896)
│    │    │    └─Upsample: 4-12                         [1, 128, 40, 40]          --
│    │    │    └─Concat: 4-13                           [1, 256, 40, 40]          --
│    │    │    └─C3: 4-14                               [1, 128, 40, 40]          (90,496)
│    │    │    └─Conv: 4-15                             [1, 64, 40, 40]           (8,256)
│    │    │    └─Upsample: 4-16                         [1, 64, 80, 80]           --
│    │    │    └─Concat: 4-17                           [1, 128, 80, 80]          --
│    │    │    └─C3: 4-18                               [1, 64, 80, 80]           (22,720)
│    │    │    └─Conv: 4-19                             [1, 64, 40, 40]           (36,928)
│    │    │    └─Concat: 4-20                           [1, 128, 40, 40]          --
│    │    │    └─C3: 4-21                               [1, 128, 40, 40]          (74,112)
│    │    │    └─Conv: 4-22                             [1, 128, 20, 20]          (147,584)
│    │    │    └─Concat: 4-23                           [1, 256, 20, 20]          --
│    │    │    └─C3: 4-24                               [1, 256, 20, 20]          (295,680)
│    │    │    └─Detect: 4-25                           [1, 25200, 85]            (115,005)
=========================================================================================================
Total params: 1,867,405
Trainable params: 0
Non-trainable params: 1,867,405
Total mult-adds (Units.GIGABYTES): 2.25
=========================================================================================================
Input size (MB): 4.92
Forward/backward pass size (MB): 111.75
Params size (MB): 7.47
Estimated Total Size (MB): 124.14
=========================================================================================================

But since YOLOv6, it seems that Ultralytics forces us to load models with their package:

# From their doc
from ultralytics import YOLO
model = YOLO("yolov10n.pt")  # Load pre-trained model

And now, this model is callable using, for example, model(np.zeros((640, 640, 3))). But when I try using torchinfo.summary on model, I get weird things happening (I caught pieces of the console output):

>>> torchinfo.summary(model, (640,640,3))
Ultralytics YOLOv8.2.23 🚀 Python-3.12.3 torch-2.3.0+cu121 CPU (Intel Core(TM) i5-10500T 2.30GHz)
engine/trainer: task=detect, mode=train, model=yolov5n.pt, data=coco.yaml, epochs=100, time=None, patience=100, batch=16, imgsz=640, save=True, save_period=-1, cache=False, device=None, workers=8, project=None, name=train14, exist_ok=False, pretrained=True, optimizer=auto, verbose=True, seed=0, deterministic=True, single_cls=False, rect=False, cos_lr=False, close_mosaic=10, resume=False, amp=True, fraction=1.0, profile=False, freeze=None, multi_scale=False, overlap_mask=True, mask_ratio=4, dropout=0.0, val=True, split=val, save_json=False, save_hybrid=False, conf=None, iou=0.7, max_det=300, half=False, dnn=False, plots=True, source=None, vid_stride=1, stream_buffer=False, visualize=False, augment=False, agnostic_nms=False, classes=None, retina_masks=False, embed=None, show=False, save_frames=False, save_txt=False, save_conf=False, save_crop=False, show_labels=True, show_conf=True, show_boxes=True, line_width=None, format=torchscript, keras=False, optimize=False, int8=False, dynamic=False, simplify=False, opset=None, workspace=4, nms=False, lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, dfl=1.5, pose=12.0, kobj=1.0, label_smoothing=0.0, nbs=64, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, bgr=0.0, mosaic=1.0, mixup=0.0, copy_paste=0.0, auto_augment=randaugment, erasing=0.4, crop_fraction=1.0, cfg=None, tracker=botsort.yaml, save_dir=[...]/.pyenv/runs/detect/train14

                   from  n    params  module                                       arguments                     
  0                  -1  1      1760  ultralytics.nn.modules.conv.Conv             [3, 16, 6, 2, 2]              
  1                  -1  1      4672  ultralytics.nn.modules.conv.Conv             [16, 32, 3, 2]                
  2                  -1  1      4800  ultralytics.nn.modules.block.C3              [32, 32, 1]                   
  3                  -1  1     18560  ultralytics.nn.modules.conv.Conv             [32, 64, 3, 2]                
  4                  -1  2     29184  ultralytics.nn.modules.block.C3              [64, 64, 2]                   
  5                  -1  1     73984  ultralytics.nn.modules.conv.Conv             [64, 128, 3, 2]               
  6                  -1  3    156928  ultralytics.nn.modules.block.C3              [128, 128, 3]                 
  7                  -1  1    295424  ultralytics.nn.modules.conv.Conv             [128, 256, 3, 2]              
  8                  -1  1    296448  ultralytics.nn.modules.block.C3              [256, 256, 1]                 
  9                  -1  1    164608  ultralytics.nn.modules.block.SPPF            [256, 256, 5]                 
 10                  -1  1     33024  ultralytics.nn.modules.conv.Conv             [256, 128, 1, 1]              
 11                  -1  1         0  torch.nn.modules.upsampling.Upsample         [None, 2, 'nearest']          
 12             [-1, 6]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 13                  -1  1     90880  ultralytics.nn.modules.block.C3              [256, 128, 1, False]          
 14                  -1  1      8320  ultralytics.nn.modules.conv.Conv             [128, 64, 1, 1]               
 15                  -1  1         0  torch.nn.modules.upsampling.Upsample         [None, 2, 'nearest']          
 16             [-1, 4]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 17                  -1  1     22912  ultralytics.nn.modules.block.C3              [128, 64, 1, False]           
 18                  -1  1     36992  ultralytics.nn.modules.conv.Conv             [64, 64, 3, 2]                
 19            [-1, 14]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 20                  -1  1     74496  ultralytics.nn.modules.block.C3              [128, 128, 1, False]          
 21                  -1  1    147712  ultralytics.nn.modules.conv.Conv             [128, 128, 3, 2]              
 22            [-1, 10]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 23                  -1  1    296448  ultralytics.nn.modules.block.C3              [256, 256, 1, False]          
 24        [17, 20, 23]  1    897664  ultralytics.nn.modules.head.Detect           [80, [64, 128, 256]]          
YOLOv5n summary: 262 layers, 2654816 parameters, 2654800 gradients, 7.8 GFLOPs

Transferred 427/427 items from pretrained weights
Freezing layer 'model.24.dfl.conv.weight'
train: WARNING ⚠️ [...]/datasets/coco/images/train2017/000000118590.jpg: ignoring corrupt image/label: [Errno 2] No such file or directory: '[...]/datasets/coco/images/train2017/000000118590.jpg'
train: WARNING ⚠️ [...]/datasets/coco/images/train2017/000000118598.jpg: ignoring corrupt image/label: [Errno 2] No such file or directory: '[...]/datasets/coco/images/train2017/000000118598.jpg'
train: WARNING ⚠️ [...]/datasets/coco/images/train2017/000000118605.jpg: ignoring corrupt image/label: [Errno 2] No such file or directory: '[...]/datasets/coco/images/train2017/000000118605.jpg'
train: WARNING ⚠️ [...]/datasets/coco/images/train2017/000000118606.jpg: ignoring corrupt image/label: [Errno 2] No such file or directory: '[...]/datasets/coco/images/train2017/000000118606.jpg'
train: WARNING ⚠️ [...]/datasets/coco/images/train2017/000000118607.jpg: ignoring corrupt image/label: [Errno 2] No such file or directory: '[...]/datasets/coco/images/train2017/000000118607.jpg'
train: WARNING ⚠️ [...]/datasets/coco/images/train2017/000000118612.jpg: ignoring corrupt image/label: [Errno 2] No such file or directory: '[...]/datasets/coco/images/train2017/000000118612.jpg'
train: WARNING ⚠️ [...]/datasets/coco/images/train2017/000000118614.jpg: ignoring corrupt image/label: [Errno 2] No such file or directory: '[...]/datasets/coco/images/train2017/000000118614.jpg'
train: WARNING ⚠️ [...]/datasets/coco/images/train2017/000000118615.jpg: ignoring corrupt image/label: [Errno 2] No such file or directory: '[...]/datasets/coco/images/train2017/000000118615.jpg'
Traceback (most recent call last):
  File "[...]/lib/python3.12/site-packages/torchinfo/torchinfo.py", line 286, in forward_pass
    model.eval()
  File "[...]/lib/python3.12/site-packages/torch/nn/modules/module.py", line 2449, in eval
    return self.train(False)
           ^^^^^^^^^^^^^^^^^
  File "[...]/lib/python3.12/site-packages/ultralytics/engine/model.py", line 674, in train
    self.trainer.train()
  File "[...]/lib/python3.12/site-packages/ultralytics/engine/trainer.py", line 199, in train
    self._do_train(world_size)
  File "[...]/lib/python3.12/site-packages/ultralytics/engine/trainer.py", line 313, in _do_train
    self._setup_train(world_size)
  File "[...]/lib/python3.12/site-packages/ultralytics/engine/trainer.py", line 277, in _setup_train
    self.train_loader = self.get_dataloader(self.trainset, batch_size=batch_size, rank=RANK, mode="train")
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/lib/python3.12/site-packages/ultralytics/models/yolo/detect/train.py", line 49, in get_dataloader
    dataset = self.build_dataset(dataset_path, mode, batch_size)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/lib/python3.12/site-packages/ultralytics/models/yolo/detect/train.py", line 43, in build_dataset
    return build_yolo_dataset(self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[...]/lib/python3.12/site-packages/ultralytics/data/build.py", line 87, in build_yolo_dataset
    return dataset(
           ^^^^^^^^
  File "[...]/lib/python3.12/site-packages/ultralytics/data/dataset.py", line 64, in __init__
    super().__init__(*args, **kwargs)
  File "[...]/lib/python3.12/site-packages/ultralytics/data/base.py", line 74, in __init__
    self.labels = self.get_labels()
                  ^^^^^^^^^^^^^^^^^
  File "[...]/lib/python3.12/site-packages/ultralytics/data/dataset.py", line 161, in get_labels
    len_cls, len_boxes, len_segments = (sum(x) for x in zip(*lengths))
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: not enough values to unpack (expected 3, got 0)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "[...]/lib/python3.12/site-packages/torchinfo/torchinfo.py", line 304, in forward_pass
    raise RuntimeError(
RuntimeError: Failed to run torchinfo. See above stack traces for more details. Executed layers up to: [SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7, SiLU: 7]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "[...]/lib/python3.12/site-packages/torchinfo/torchinfo.py", line 223, in summary
    summary_list = forward_pass(
                   ^^^^^^^^^^^^^
  File "[...]/lib/python3.12/site-packages/torchinfo/torchinfo.py", line 313, in forward_pass
    model.train(saved_model_mode)
  File "[...]/lib/python3.12/site-packages/ultralytics/engine/model.py", line 655, in train
    self.trainer = (trainer or self._smart_load("trainer"))(overrides=args, _callbacks=self.callbacks)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: 'bool' object is not callable
>>> 

I don't exactly know what torchinfo.summary does under the hood to call the model, but it seems to be failing. Any tips? I trying playing with a batch_dim, to no avail.

Thanks in advance.
Cheers!