Input size disappears between torch.jit.save and torch.jit load

elenkalda-arm: <p>Hello! Apologies if this is a silly question, I’m very new to PyTorch - basically, I’m trying to save a traced model with <code>torch.jit.save</code> and then later load it with <code>torch.jit.load</code> and access the size of the input tensor, but it looks like the size information goes missing between saving and loading. If I run the snippet</p>
<pre><code class="lang-auto">import torch
import torchvision
model_name = "resnet18"

model = getattr(torchvision.models, model_name)(pretrained=True)
model = model.eval()

input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)
traced_model = torch.jit.trace(model, input_data)

print(list(traced_model.graph.inputs()))
torch.jit.save(traced_model, "resnet18.pth")

loaded_model = torch.jit.load("resnet18.pth")
print(list(loaded_model.graph.inputs()))
</code></pre>
<p>the first print (corresponding to the model before saving) prints out</p>
<pre><code class="lang-auto">[self.1 defined in (%self.1 : __torch__.torchvision.models.resnet.ResNet, %input.1 : Float(1:150528, 3:50176, 224:224, 224:1, requires_grad=0, device=cpu) = prim::Param()
), input.1 defined in (%self.1 : __torch__.torchvision.models.resnet.ResNet, %input.1 : Float(1:150528, 3:50176, 224:224, 224:1, requires_grad=0, device=cpu) = prim::Param()
)]
</code></pre>
<p>which shows some familiar numbers about the inputs, but the second print (corresponding to the loaded model) prints</p>
<pre><code class="lang-auto">[self.1 defined in (%self.1 : __torch__.torchvision.models.resnet.ResNet, %input.1 : Tensor = prim::Param()
), input.1 defined in (%self.1 : __torch__.torchvision.models.resnet.ResNet, %input.1 : Tensor = prim::Param()
)]
</code></pre>
<p>The input numbers have gone missing… Does anyone know why this happens? Or maybe there is some other way of saving a model on a disk and loading it later such that the information about inputs is still accessible? I’m asking that because I’m working on an ahead of time compiler which relies on fetching the input shape from the model.</p>

erpang007: <p>Did you find out the solution? I am struggling to find the shape for this as well.</p>

Ashvant_Ram_Selvam: <p>Were you able to get to the solution? Would appreciate if shared.</p>

fabianpedd: <p>It’s not a direct solution, but you could export the model to ONNX and use that instead:</p>
<pre><code class="lang-auto">import torch
import torchvision
import onnx

model_name = "resnet18"

model = getattr(torchvision.models, model_name)(pretrained=True)
model = model.eval()

input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)

# Export model using onnx
torch.onnx.export(
    model,
    input_data,
    "resnet18.onnx",
    input_names=["input"],
    output_names=["output"],
    verbose=False,
)

# Load model again
onnx_model = onnx.load("resnet18.onnx")

# Check model (optional)
onnx.checker.check_model(onnx_model)

# Infer shapes of tensors inside the model (optional)
onnx_model_inferred = onnx.shape_inference.infer_shapes(onnx_model)

# Get input shape
input_dim_onnx = [d.dim_value for d in onnx_model_inferred.graph.input[0].type.tensor_type.shape.dim]
print(input_dim_onnx) # [1, 3, 224, 224]
</code></pre>