Lightning experiment

lightning_experiment.py

class BasicCNNExperiment(BaseExperiment):

    def create_model(self) -> nn.Module:
        return nn.Sequential(
            nn.Conv2d(3, 16, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(16, 32, 3),
            nn.ReLU(inplace=True),
            nn.AdaptiveMaxPool2d(1),
            nn.Flatten(1, -1),
            nn.Linear(32, self.hparams.n_classes)
        )

    def forward(self, batch: Dict) -> Any:
        features, targets = batch
        out = self.model(features)
        loss = cross_entropy(out, targets)
        return {'loss': loss, 'outputs': out}