Untitled

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

# Define a deep convolutional neural network without MaxPool2d
class DeepConvNetNoPooling(nn.Module):
    def __init__(self, in_channels=3, num_classes=10):
        super(DeepConvNetNoPooling, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1),  # Layer 1
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),           # Layer 2 (stride reduces spatial size)
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),          # Layer 3 (stride reduces spatial size)
            nn.ReLU(),
            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),          # Layer 4
            nn.ReLU(),
            nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),          # Layer 5 (stride reduces spatial size)
            nn.ReLU()
        )
        self.classifier = nn.Sequential(
            nn.Linear(512 * 8 * 8, 1024),                                     # Adjust input size based on final feature map
            nn.ReLU(),
            nn.Linear(1024, num_classes)                                      # Output layer
        )

    def forward(self, x):
        x = self.features(x)
        # Replace Flatten with Reshape
        x = x.view(x.size(0), -1)  # Reshape tensor to (batch_size, -1)
        x = self.classifier(x)
        return x

# Initialize the model
model = DeepConvNetNoPooling(in_channels=3, num_classes=10)
model.eval()  # Set 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 = "deep_conv_net_no_pooling.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"Deep CNN model without MaxPool2d has been converted to ONNX and saved to {onnx_path}")