#7520 Do Upsample with ConvTranspose2d

Search before asking

I have searched the YOLOv5 issues and found no similar feature requests.

Description

The origin nn.Upsample layer may not be supported on some edge devices like Hi3559A.

I have trained YOLOv5s with nn.ConvTranspose2d which is compatible with Caffe 1.0 and Hi3559A's NNIE (#5397). The training process depends on 4xV100 GPU and I used the same config as the repo's pretrained checkpoints to validate the models.

# train
python -m torch.distributed.run --nproc_per_node 4 train.py --device 0,1,2,3 --data data/coco.yaml --hyp data/hyps/hyp.scratch-low.yaml  --cfg path/to/model.yaml --batch 512 --epochs 300 --weights ''
# val
python val.py --verbose --data data/coco.yaml --conf 0.001 --iou 0.65 --batch 1 --weights path/to/model.pt

Here is the val result on COCO val2017 dataset. It seems that with larger deconv kernel, the model reaches better performance. I only compare kernel size 2x2 and 4x4 while the model with 4x4 deconv kernel got the best performance.

Model	deconv kernel size	size (pixels)	mAP @0.5:0.95	mAP @0.5	Speed V100 b1(ms)	Speed V100 b32(ms)	params (M)	FLOPs @640(B)
YOLOv5s	-	640	33.7	52.9	5.6	2.2	7.23	16.5
YOLOv5s-deconv	2	640	33.4	52.5	5.6	2.4	7.55	18.2
YOLOv5s-deconv	4	640	34.7	54.2	5.8	2.5	8.54	23.2

I hope you can test it under your environment for fair comparison and if it works, it couldn't be better to add the yolov5-deconv to the model hub.

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']], # the origin nn.Upsample layer
   # [-1, 1, nn.ConvTranspose2d, [512, 2, 2]], # nn.ConvTransposed2d with 2x2 deconv kernel
   # [-1, 1, nn.ConvTranspose2d, [512, 4, 2, 1]], # nn.ConvTransposed2d with 4x4 deconv kernel
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']], # the origin nn.Upsample layer
   # [-1, 1, nn.ConvTranspose2d, [256, 2, 2]], # nn.ConvTransposed2d with 2x2 deconv kernel
   # [-1, 1, nn.ConvTranspose2d, [256, 4, 2, 1]], # nn.ConvTransposed2d with 4x4 deconv kernel
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

# the parse_model of yolo.py should be changed from
# if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
#         BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]:
# to
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
         BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d]:

Use case

The yolov5-deconv is compatible with Caffe 1.0 and the NNIE of Hi3559A, making it easier to deploy the model on some edge devices.

Additional

I have not test the nn.ConvTranspose2d under larger model like yolov5l. To get better performance, the deconv kernel size may be even larger in large models since there are more features during large model's forward process.

Are you willing to submit a PR?

Yes I'd like to help by submitting a PR!