Untitled

import torch
import torch.nn as nn
import torch.onnx

# Define a single-layer convolutional model
class SingleLayerConvNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=16, kernel_size=3):
        super(SingleLayerConvNet, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=1)

    def forward(self, x):
        return self.conv(x)

# Initialize the model
model = SingleLayerConvNet(in_channels=3, out_channels=16, kernel_size=3)
model.eval()  # Set the model to evaluation mode

# Dummy input for ONNX export (batch size 1, 3 channels, 64x64 image)
dummy_input = torch.randn(1, 3, 64, 64)

# Path to save the ONNX model
onnx_path = "single_layer_conv.onnx"

# Export the model to ONNX
torch.onnx.export(
    model,                  # Model to export
    dummy_input,            # Example input
    onnx_path,              # Path to save the ONNX file
    export_params=True,     # Store trained parameter weights inside the model
    opset_version=11,       # ONNX opset version
    input_names=["input"],  # Input tensor name
    output_names=["output"] # Output tensor name
)

print(f"Model has been converted to ONNX and saved to {onnx_path}")