1

I am trying to compile my PyTorch model into mlir using torch-mlir. The compile method from torch-mlir requires three arguments: 1) model to convert 2) example arguments to use when inferring the shapes of the arguments to the forward method of the model 3) desired output type.

My question is with the second input. I am not sure what the shape is of the arguments to the forward method. The model I am using to practice with is the Iris dataset. The full code is here.

The model and forward method are:

class Model(nn.Module):
    def __init__(self, input_dim):
        super(Model, self).__init__()
        self.layer1 = nn.Linear(input_dim, 50)
        self.layer2 = nn.Linear(50, 3)
        
    def forward(self, x):
        print(x.shape)
        x = F.relu(self.layer1(x))
        x = F.softmax(self.layer2(x), dim=1)
        return x
RuthC
  • 1,269
  • 1
  • 10
  • 21
PrematureCorn
  • 121
  • 5

1 Answers1

0

From the repository you provided:

example_args: A list of example arguments to use when inferring the shapes of the arguments to forward method of the model. A single tensor is treated as a list of a single tensor. A TensorPlaceholder object is also allowed in the place of any Tensor. For models with multiple methods, an ExampleArgs object can be passed.

This means that if x has a shape of (sample_size, input_dims), the argument for example_args should be (number_of_examples, sample_size, input_dims).

Quantum
  • 510
  • 1
  • 2
  • 19
  • 1
    Thanks for the response Quantum! So if the training data is a 150x4 array then `sample size = 150` and `input_dims = 4` if I'm understand correctly. I guess the last thing I'm confused about is the number of examples. I think just setting this to one may be fine? – PrematureCorn Jun 06 '23 at 13:48
  • Exactly, this is how I understood the documentation. Based on your data layout, the sample size and input_dims could be reversed, but the "normal" approach is to have an array of samples that each contain an array of inputs. The number of examples is not specified, so it is safe to assume that any shape works well. The easiest way would be to include only one example. – Quantum Jun 06 '23 at 13:51