syheliel opened a new issue, #17879:
URL: https://github.com/apache/tvm/issues/17879
### Environment
- torch 1.13.1
- apache-tvm 0.14.dev273
### Steps to reproduce
```
import torch
import tvm
from tvm import relay
import torch.nn as nn
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
# First convolutional layer
self.conv1 = nn.Conv2d(in_channels=3, out_channels=32,
kernel_size=3, padding=1)
self.relu1 = nn.ReLU()
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
# Second convolutional layer
self.conv2 = nn.Conv2d(in_channels=32, out_channels=64,
kernel_size=3, padding=1)
self.relu2 = nn.ReLU()
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
# Third convolutional layer
self.conv3 = nn.Conv2d(in_channels=64, out_channels=128,
kernel_size=3, padding=1)
self.relu3 = nn.ReLU()
self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
# Fully connected layers
self.fc1 = nn.Linear(128 * 4 * 4, 512)
self.relu4 = nn.ReLU()
self.dropout = nn.Dropout(0.5)
self.fc2 = nn.Linear(512, 10) # 10 classes for classification
def forward(self, x):
# First conv block
x = self.conv1(x)
x = self.relu1(x)
x = self.pool1(x)
# Second conv block
x = self.conv2(x)
x = self.relu2(x)
x = self.pool2(x)
# Third conv block
x = self.conv3(x)
x = self.relu3(x)
x = self.pool3(x)
# Flatten the output
x = x.view(x.size(0), -1)
# Fully connected layers
x = self.fc1(x)
x = self.relu4(x)
x = self.dropout(x)
x = self.fc2(x)
return x
# Example usage
if __name__ == "__main__":
# Create model instance
model = SimpleCNN()
# Create a sample input tensor (batch_size=1, channels=3, height=32,
width=32)
sample_input = torch.randn(1, 3, 32, 32)
scripted_model = torch.jit.script(model)
shape_list = ["sample_input", sample_input.shape]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]