How to export video classification model to torch script

jit
1

pytorch version : 1.7.1
os : win10 64

Trying to export the video classification by following script

import os
import numpy as np
import torch
import torchvision

model = torchvision.models.video.r3d_18(pretrained=True, progress=True)
model.eval()

img = torch.zeros((16, 3, 112, 112))
_ = model(img)  # dry run
traced_script_module = torch.jit.trace(model, img)
traced_script_module.save("video_r3d_18.pt")

The script give me error messages

RuntimeError: Expected 5-dimensional input for 5-dimensional weight [64, 3, 3, 7, 7], but got 4-dimensional input of size [3, 16, 112, 112] instead

Afte I change it to [64, 3, 3, 7, 7], I can export the model, but the training codes and the doc, both use [3,16,112,112], this is weird, why this happen?

Thanks

2

The error is raised in the “dry run” section, bot the tracing, since you are passing a 4-dimensional input, while 3D models expect a 5-dimensional input.
You could unsqueeze the “depth” dimension and the code should work:

img = torch.zeros((16, 3, 1, 112, 112))
_ = model(img)  # dry run
1 Like