Search code examples
pythonmachine-learningpytorch

Create custom model class with python


I have an example class as follows:

class MLP(nn.Module):
    # Declare a layer with model parameters. Here, we declare two fully
    # connected layers
    def __init__(self):
        # Call the constructor of the `MLP` parent class `Module` to perform
        # the necessary initialization. In this way, other function arguments
        # can also be specified during class instantiation, such as the model
        # parameters, `params` (to be described later)
        super().__init__()
        self.hidden = nn.Linear(20, 256)  # Hidden layer
        self.out = nn.Linear(256, 10)  # Output layer

    # Define the forward propagation of the model, that is, how to return the
    # required model output based on the input `X`
    def forward(self, X):
        # Note here we use the funtional version of ReLU defined in the
        # nn.functional module.
        return self.out(torch.relu(self.hidden(X)))

Calling the class is like this: net = MLP() net(X)

Now, I need to create a similar class and function for a model with 4 layers:

Layers Configuration Activation Function fully connected input size 128, output size 64 ReLU fully connected input size 64, output size 32 ReLU dropout probability 0.5 - fully connected input size 32, output size 1 Sigmoid

I need to pass the following assertion:

model = Net()

assert model.fc1.in_features == 128
assert model.fc1.out_features == 64
assert model.fc2.in_features == 64
assert model.fc2.out_features == 32
assert model.fc3.in_features == 32
assert model.fc3.out_features == 1

x = torch.rand(2, 128)
output = model.forward(x)
assert output.shape == (2, 1), "Net() is wrong!"

Here is what I have so far:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
                
        
        self.fc1 = nn.Linear(128, 64)
        self.fc2 = nn.Linear(64, 32)
        self.dropout = nn.Dropout(p=0.5)
        self.fc3 = nn.Linear(32, 1)
        

    def forward(self, x):
        return self.fc3(torch.sigmoid(self.dropout(self.fc2(torch.relu(self.fc1(torch.relu(X)))))))
       

But I'm getting an error:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x20 and 128x64)

How to resolve it?


Solution

  • You have a capital X in the forward method, while the function expects lowercase x. You must have a tensor with shape (2, 20) assigned to variable capital X.