export model to onnx error.

[morestart]

/home/cat/PycharmProjects/YOLOM/yolo/detectHead/detect.py:39: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
Traceback (most recent call last):
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/torch/onnx/utils.py", line 729, in _export
    opened_file.write(proto)
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/torch/nn/modules/module.py", line 947, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'YOLO' object has no attribute 'write'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/cat/PycharmProjects/YOLOM/export.py", line 90, in <module>
    run(
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/cat/PycharmProjects/YOLOM/export.py", line 86, in run
    export_onnx(model, im, Path(model_path), opset, False, dynamic, simplify)
  File "/home/cat/PycharmProjects/YOLOM/export.py", line 19, in export_onnx
    model.to_onnx(model, im, opset_version=opset,
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1899, in to_onnx
    torch.onnx.export(self, input_sample, file_path, **kwargs)
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/torch/onnx/__init__.py", line 271, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/torch/onnx/utils.py", line 88, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/torch/onnx/utils.py", line 729, in _export
    opened_file.write(proto)
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/torch/serialization.py", line 225, in __exit__
    self.file_like.flush()
  File "/home/cat/miniconda3/envs/YOLOM/lib/python3.8/site-packages/torch/nn/modules/module.py", line 947, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'YOLO' object has no attribute 'flush'

this is my code:

def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix='ONNX:'):
    # YOLOv5 ONNX export
    # try:
        # check_requirements(('onnx',))
    import onnx

    print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
    f = file.with_suffix('.onnx')

    model.to_onnx(model, im, opset_version=opset,
                  training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
                  do_constant_folding=not train,
                  input_names=['images'],
                  output_names=['output'],
                  dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)
                                'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
                                } if dynamic else None)

    # Checks
    model_onnx = onnx.load(f)  # load onnx model
    onnx.checker.check_model(model_onnx)  # check onnx model
    # print(onnx.helper.printable_graph(model_onnx.graph))  # print

    # Simplify
    if simplify:
        try:
            # check_requirements(('onnx-simplifier',))
            import onnxsim

            print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
            model_onnx, check = onnxsim.simplify(
                model_onnx,
                dynamic_input_shape=dynamic,
                input_shapes={'images': list(im.shape)} if dynamic else None)
            assert check, 'assert check failed'
            onnx.save(model_onnx, f)
        except Exception as e:
            print(f'{prefix} simplifier failure: {e}')
    print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
    print(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
    # except Exception as e:
    #     print(f'{prefix} export failure: {e}')


@torch.no_grad()
def run(
        model_path: str,
        half: bool = False,
        imgsz=(640, 640),
        device: str = 'cuda:0',
        batch_size: int = 1,
        dynamic=False,
        simplify=False,
        inplace=False,
        opset=12,
):
    model = YOLO.load_from_checkpoint(model_path)
    model.to(device)
    model.eval()

    imgsz *= 2 if len(imgsz) == 1 else 1
    gs = int(max(model.stride))  # grid size (max stride)
    imgsz = [check_img_size(x, gs) for x in imgsz]  # verify img_size are gs-multiples
    im = torch.zeros(batch_size, 3, *imgsz).to(device)
    if half:
        assert not (device == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
        im, model = im.half(), model.half()

    for k, m in model.named_modules():
        # if isinstance(m, Conv):  # assign export-friendly activations
        #     if isinstance(m.act, nn.SiLU):
        #         m.act = SiLU()
        if isinstance(m, Detect):
            m.inplace = inplace
            m.onnx_dynamic = dynamic

    export_onnx(model, im, Path(model_path), opset, False, dynamic, simplify)


if __name__ == '__main__':
    run(
        model_path='cspdarknet53-epoch=15.ckpt',
    )

[rohitgr7]

hey @morestart !

model isn't the first argument. It should be just model.to_onnx(im, ...)