MegEngine implementation of YOLOX

When i used this repo on MegStudio and tried to train yolox_tiny with the pretrained model, an error occurred. The detail log are as follow.

2021-09-15 13:11:11 | INFO | yolox.core.trainer:247 - loading checkpoint for fine tuning 2021-09-15 13:11:11 | ERROR | main:93 - An error has been caught in function '', process 'MainProcess' (359), thread 'MainThread' (139974572922688): Traceback (most recent call last):

File "tools/train.py", line 93, in main(exp, args) │ │ └ Namespace(batch_size=16, ckpt='yolox_tiny.pkl', devices=1, exp_file='exps/default/yolox_tiny.py', experiment_name='yolox_tiny... │ └ ╒══════════════════╤═════════════════════════════════════════════════════════════════════════════════════════════════════════... └ <function main at 0x7f4e5d7308c0>

File "tools/train.py", line 73, in main trainer.train() │ └ <function Trainer.train at 0x7f4dec68b680> └ <yolox.core.trainer.Trainer object at 0x7f4d9a68a7d0>

File "/home/megstudio/workspace/YOLOX/yolox/core/trainer.py", line 46, in train self.before_train() │ └ <function Trainer.before_train at 0x7f4d9a6f55f0> └ <yolox.core.trainer.Trainer object at 0x7f4d9a68a7d0>

File "/home/megstudio/workspace/YOLOX/yolox/core/trainer.py", line 107, in before_train model = self.resume_train(model) │ │ └ YOLOX( │ │ (backbone): YOLOPAFPN( │ │ (backbone): CSPDarknet( │ │ (stem): Focus( │ │ (conv): BaseConv( │ │ (conv): ... │ └ <function Trainer.resume_train at 0x7f4d9a70c0e0> └ <yolox.core.trainer.Trainer object at 0x7f4d9a68a7d0>

File "/home/megstudio/workspace/YOLOX/yolox/core/trainer.py", line 249, in resume_train ckpt = mge.load(ckpt_file, map_location="cpu")["model"] │ │ └ 'yolox_tiny.pkl' │ └ <function load at 0x7f4df6c46680> └ <module 'megengine' from '/home/megstudio/.miniconda/envs/xuan/lib/python3.7/site-packages/megengine/init.py'>

KeyError: 'model'

 python tools/demo.py image -n yolox-tiny -c /path/to/your/yolox_tiny.pkl --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 416 --save_result --device gpu
2021-09-07 18:45:49.600 | INFO     | __main__:main:250 - Args: Namespace(camid=0, ckpt='/path/to/your/yolox_tiny.pkl', conf=0.25, demo='image', device='gpu', exp_file=None, experiment_name='yolox_tiny', fp16=False, fuse=False, legacy=False, name='yolox-tiny', nms=0.45, path='assets/dog.jpg', save_result=True, trt=False, tsize=416)
E:\anaconda3\envs\YOLOX\lib\site-packages\torch\nn\functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  ..\c10/core/TensorImpl.h:1156.)
  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
2021-09-07 18:45:49.791 | INFO     | __main__:main:260 - Model Summary: Params: 5.06M, Gflops: 6.45
Traceback (most recent call last):
  File "tools/demo.py", line 306, in <module>
    main(exp, args)
  File "tools/demo.py", line 263, in main
    model.cuda()
  File "E:\anaconda3\envs\YOLOX\lib\site-packages\torch\nn\modules\module.py", line 637, in cuda
    return self._apply(lambda t: t.cuda(device))
  File "E:\anaconda3\envs\YOLOX\lib\site-packages\torch\nn\modules\module.py", line 530, in _apply
    module._apply(fn)
  File "E:\anaconda3\envs\YOLOX\lib\site-packages\torch\nn\modules\module.py", line 530, in _apply
    module._apply(fn)
  File "E:\anaconda3\envs\YOLOX\lib\site-packages\torch\nn\modules\module.py", line 530, in _apply
    module._apply(fn)
  [Previous line repeated 2 more times]
  File "E:\anaconda3\envs\YOLOX\lib\site-packages\torch\nn\modules\module.py", line 552, in _apply
    param_applied = fn(param)
  File "E:\anaconda3\envs\YOLOX\lib\site-packages\torch\nn\modules\module.py", line 637, in <lambda>
    return self._apply(lambda t: t.cuda(device))
  File "E:\anaconda3\envs\YOLOX\lib\site-packages\torch\cuda\__init__.py", line 166, in _lazy_init
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled

环境 CUDA Version: 11.2 没问题

按照官方的教程 报错

Testing Devices

x86_64 -- Intel(R) Xeon(R) CPU E5-2620 v4 @ 2.10GHz AArch64 -- xiamo phone mi9 CUDA -- 1080TI @ cuda-10.1-cudnn-v7.6.3-TensorRT-6.0.1.5.sh @ Intel(R) Xeon(R) CPU E5-2620 v4 @ 2.10GHz

Shouldn't it be Xiaomi phone mi9?

1. img_info for VOC dataset is wrong.
2. grid for yolo_head is wrong (Similar to https://github.com/MegEngine/YOLOX/issues/9). If the image has the same height and width, it will be ok. But, when height != weight, it will be wrong.

当输入宽高不一致时报错， 在训练过程中报错，报错时机随缘： yolo_head.py", line 351, in get_assignments bboxes_preds_per_image = bboxes_preds_per_image[fg_mask] RuntimeError: assertion `dtype == dst.dtype && dst.is_contiguous()' failed at ../../../../../../dnn/src/common/elemwise/opr_impl.cpp:281: void megdnn::ElemwiseForward::check_layout_and_broadcast(const TensorLayoutPtrArray&, const megdnn::TensorLayout&)