(self)
| 48 | |
| 49 | class MnistModel(LightningModule): |
| 50 | def __init__(self): |
| 51 | super().__init__() |
| 52 | |
| 53 | # Tunable parameters |
| 54 | self.hidden_size_1 = args.hidden_size_1 |
| 55 | self.hidden_size_2 = args.hidden_size_2 |
| 56 | self.learning_rate = args.learning_rate |
| 57 | self.dropout = args.dropout |
| 58 | self.batch_size = args.batch_size |
| 59 | |
| 60 | # Set class attributes |
| 61 | self.data_dir = PATH_DATASETS |
| 62 | |
| 63 | # Hardcode some dataset specific attributes |
| 64 | self.num_classes = 10 |
| 65 | self.dims = (1, 28, 28) |
| 66 | channels, width, height = self.dims |
| 67 | self.transform = transforms.Compose( |
| 68 | [ |
| 69 | transforms.ToTensor(), |
| 70 | transforms.Normalize((0.1307,), (0.3081,)), |
| 71 | ] |
| 72 | ) |
| 73 | |
| 74 | # Create a PyTorch model |
| 75 | layers = [nn.Flatten()] |
| 76 | width = channels * width * height |
| 77 | hidden_layers = [self.hidden_size_1, self.hidden_size_2] |
| 78 | num_params = 0 |
| 79 | for hidden_size in hidden_layers: |
| 80 | if hidden_size > 0: |
| 81 | layers.append(nn.Linear(width, hidden_size)) |
| 82 | layers.append(nn.ReLU()) |
| 83 | layers.append(nn.Dropout(self.dropout)) |
| 84 | num_params += width * hidden_size |
| 85 | width = hidden_size |
| 86 | layers.append(nn.Linear(width, self.num_classes)) |
| 87 | num_params += width * self.num_classes |
| 88 | |
| 89 | # Save the model and parameter counts |
| 90 | self.num_params = num_params |
| 91 | self.model = nn.Sequential(*layers) # No need to use Relu for the last layer |
| 92 | |
| 93 | def forward(self, x): |
| 94 | x = self.model(x) |
nothing calls this directly
no outgoing calls
no test coverage detected