I am learning to convert models to the onnx. I wrote a simple classification model to classify cats and dogs but the onnx conversion is giving shape unmatch error.
vision
class Cnn(nn.Module):
def init(self):
super(Cnn,self).init()
self.layer1 = nn.Sequential(
nn.Conv2d(3,16,kernel_size=3, padding=0,stride=2),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.layer2 = nn.Sequential(
nn.Conv2d(16,32, kernel_size=3, padding=0, stride=2),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.layer3 = nn.Sequential(
nn.Conv2d(32,64, kernel_size=3, padding=0, stride=2),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.fc1 = nn.Linear(3*3*64,10)
self.fc2 = nn.Linear(10,2)
self.relu = nn.ReLU()
def forward(self,x):
out = self.layer1(x)
out = self.layer2(out)
out = self.layer3(out)
out = out.view(out.size(0),-1)
out = self.relu(self.fc1(out))
out = self.fc2(out)
return out
model = Cnn()
model.load_state_dict(torch.load(‘cat_dog.pth’))
dummy_input = torch.randn(1, 3, 224, 224) # Example input with shape (batch_size, channels, height, width)
torch_out = model(dummy_input)
input_names = [“input”]
output_names = [“output”]
onnx_path = “cnn_model.onnx” # Path to save the ONNX model file
Export the model to ONNX
torch.onnx.export(
model,
dummy_input,
onnx_path,
input_names=input_names,
output_names=output_names,
opset_version=10, # Use ONNX opset version 10
do_constant_folding=True,
dynamic_axes={
“input”: {0: “batch_size”},
“output”: {0: “batch_size”}
}
)
ERROR
The error says as follows: RuntimeError: Given groups=1, weight of size [16, 3, 3, 3], expected input[1, 1, 224, 224] to have 3 channels, but got 1 channels instead