search

hzwer / Practical-RIFE

We are developing more practical frame interpolation approach.
MIT License
542 stars 60 forks source link

ONNX export script #60

Closed Djdefrag closed 5 months ago

Djdefrag commented 6 months ago

Hi! thank you for your hard work.

I am trying to convert RIFE 4.13 in onnx, but without success. Maibe can you help me?

    dummy_input = torch_randn(1, 3, 256, 256, requires_grad=False)
    dymmy2      = torch_randn(1, 3, 256, 256, requires_grad=False)
    imgs        = torch_cat((dummy_input, dymmy2), 1)

    timestep = 0.5
    scale_list  = [8, 4, 2, 1]

    torch_onnx_export(
        torch_AI_model,
        (imgs, timestep, scale_list),
        f"{relative_path}{selected_AI_model}.onnx",
        verbose = True,
        do_constant_folding = True,
        export_params =True,
        input_names  = [ "input", "timestep", "scale_list" ],
        output_names = [ "flow_list", "mask", "merged" ],
        dynamic_axes = {'input': {1: 'channels', 2: 'height', 3: 'width'}}
    )
TNTwise commented 6 months ago

https://github.com/styler00dollar/VSGAN-tensorrt-docker/issues/56

Djdefrag commented 6 months ago

Thank you :D