Exercise 1: Error while building a PyTorch model by subclassing nn.Module using nn.Linear

[yasaswi-p]

Hi,
I'm trying to solve question 2 from exercise 1 which is

Build a PyTorch model by subclassing nn.Module and implement forward() function. Everything works well when I use weights and bias in the class. But when I replace them with nn.Linear(), I get an error. This also happens with the solution provided(option 2 solution).

Data:

weight = 0.3
bias = 0.9
X = torch.arange(0, 1, 0.01)
y = weight * X + bias

Data Split:

train_split = int(0.8 * len(X)) # 80% for training split
X_train, y_train = X[:train_split], y[:train_split]
X_test, y_test = X[train_split:], y[train_split:]
len(X_train), len(y_train), len(X_test), len(y_test)

Class Definition

class LRM(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_layer = nn.Linear(in_features=1,
                                     out_features=1)
    
    def forward(self, x):
        return self.linear_layer(x)

model = LRM().to(device)
optimizer = torch.optim.SGD(lr = 0.01, params = model.parameters())
loss_fn = nn.L1Loss()
epochs = 300
X_train, y_train = X_train.to(device), y_train.to(device)
X_test, y_test = X_test.to(device), y_test.to(device)

Error in the below code

for epoch in range(epochs):
  model.train()
  y_preds = model(X_train)
  loss = loss_fn(y_preds, y_train)
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()
  if epoch % 20 == 0:
    model.eval()
    with torch.inference_mode():
      test_pred = model(X_test)
      test_loss = loss_fn(test_pred, y_test)
      print(f"Epoch: {epoch} | Train loss: {loss:.3f} | Test loss: {test_loss:.3f}")

Error

RuntimeError                              Traceback (most recent call last)
Cell In [35], line 7
      5 for epoch in range(epochs):
      6   model.train()
----> 7   y_preds = model(X_train)
      8   loss = loss_fn(y_preds, y_train)
      9   optimizer.zero_grad()

File ~/mambaforge/lib/python3.10/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

Cell In [30], line 8, in LRM.forward(self, x)
      7 def forward(self, x):
----> 8     return self.linear_layer(x)

File ~/mambaforge/lib/python3.10/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/mambaforge/lib/python3.10/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)
    113 def forward(self, input: Tensor) -> Tensor:
--> 114     return F.linear(input, self.weight, self.bias)

RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x80 and 1x1)

Kindly help me with this.

[yasaswi-p]

Found the problem. It has to do with the shape of X. After adding the unsqueeze, error got resolved

weight = 0.3
bias = 0.9
X = torch.arange(0, 1, 0.01).unsqueeze(dim=1)
y = weight * X + bias

[mrdbourke]

Glad to hear you got it fixed!

[shmuhammadd]

But, please, why do we need to unsqueeze(dim=1) @mrdbourke @Pvry ?

[mrdbourke]

Hi @shmuhammad2004,

Have you tried running the code with and out without unsqueeze(dim=1) - what happens?

Does it affect the input and output shapes?