Pārlūkot izejas kodu

Tensorboard model visualization bug fix (#2758)

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:
0cae7576a9/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">
5.0
Glenn Jocher GitHub pirms 3 gadiem
vecāks
revīzija
6dd1083bbb
Šim parakstam datu bāzē netika atrasta zināma atslēga GPG atslēgas ID: 4AEE18F83AFDEB23
1 mainītis faili ar 1 papildinājumiem un 1 dzēšanām
  1. +1
    -1
      train.py

+ 1
- 1
train.py Parādīt failu

@@ -332,7 +332,7 @@ def train(hyp, opt, device, tb_writer=None):
Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
# if tb_writer:
# tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
# tb_writer.add_graph(model, imgs) # add model to tensorboard
# tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), []) # add model graph
elif plots and ni == 10 and wandb_logger.wandb:
wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
save_dir.glob('train*.jpg') if x.exists()]})

Notiek ielāde…
Atcelt
Saglabāt