Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Tensorboard model visualization bug fix #2758

Merged
merged 1 commit into from
Apr 10, 2021
Merged

Commits on Apr 10, 2021

  1. Tensorboard model visualization bug fix

    This fix should allow for visualizing YOLOv5 model graphs correctly in Tensorboard by uncommenting line 335 in train.py:
    ```python
                        if tb_writer:
                            tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), [])  # add model graph
    ```
    
    The problem was that the detect() layer checks the input size to adapt the grid if required, and tracing does not seem to like this shape check (even if the shape is fine and no grid recomputation is required). The following will warn:
    https://github.com/ultralytics/yolov5/blob/0cae7576a9241110157cd154fc2237e703c2719e/train.py#L335
    
    Solution is below. This is a YOLOv5s model displayed in TensorBoard. You can see the Detect() layer merging the 3 layers into a single output for example, and everything appears to work and visualize correctly.
    ```python
    tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), [])
    ```
    <img width="893" alt="Screenshot 2021-04-11 at 01 10 09" src="https://user-images.githubusercontent.com/26833433/114286928-349bd600-9a63-11eb-941f-7139ee6cd602.png">
    glenn-jocher committed Apr 10, 2021
    Configuration menu
    Copy the full SHA
    b1509fa View commit details
    Browse the repository at this point in the history