<jupyter_start><jupyter_text>Switch the datasets switch to Fashion Mnist<jupyter_code>def test_model(channels, width, height, hidden_size, num_classes):
    
    model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(channels * width * height, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, num_classes)
        )
    
    return model
class LitMNIST(pl.LightningModule):
    
    def __init__(self, model, data_dir='./', hidden_size=64, learning_rate=2e-4):

        super().__init__()

        # Set our init args as class attributes
        self.data_dir = data_dir
        self.hidden_size = hidden_size
        self.learning_rate = learning_rate

        # Hardcode some dataset specific attributes
        # we would need to move this out for testing different datasets
        self.num_classes = 10
        self.dims = (1, 28, 28)
        channels, width, height = self.dims
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])

        # Define PyTorch model
        self.model = model(channels, width, height, hidden_size, self.num_classes)

    def forward(self, x):
        x = self.model(x)
        return F.log_softmax(x, dim=1)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        return pl.TrainResult(loss)

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = accuracy(preds, y)
        result = pl.EvalResult(checkpoint_on=loss)

        # Calling result.log will surface up scalars for you in TensorBoard
        result.log('val_loss', loss, prog_bar=True)
        result.log('val_acc', acc, prog_bar=True)
        return result

    def test_step(self, batch, batch_idx):
        # Here we just reuse the validation_step for testing
        return self.validation_step(batch, batch_idx)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        return optimizer

    ####################
    # DATA RELATED HOOKS
    ####################

    def prepare_data(self):
        # download
        FashionMNIST(self.data_dir, train=True, download=True)
        FashionMNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):

        # Assign train/val datasets for use in dataloaders
        if stage == 'fit' or stage is None:
            FashionMnist_full = FashionMNIST(self.data_dir, train=True, transform=self.transform)
            self.Fmnist_train, self.Fmnist_val = random_split(FashionMnist_full, [55000, 5000])

        # Assign test dataset for use in dataloader(s)
        if stage == 'test' or stage is None:
            self.Fmnist_test = FashionMNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(self.Fmnist_train, batch_size=32, num_workers=8)

    def val_dataloader(self):
        return DataLoader(self.Fmnist_val, batch_size=32, num_workers=8)

    def test_dataloader(self):
        return DataLoader(self.Fmnist_test, batch_size=32, num_workers=8)
model = LitMNIST(model=test_model)
trainer = pl.Trainer(gpus=1, max_epochs=3, progress_bar_refresh_rate=20)
trainer.fit(model)
trainer.test()<jupyter_output><empty_output><jupyter_text>Trying Conv models torchvision models are built around 2xx x 2xx sized images and 3 channels<jupyter_code># Examine the resnet18
import torchvision
resnet18 = models.resnet18()
resnet18
# lets stack our own
import torch
import torch.nn
def test_model(channels, width, height, hidden_size, num_classes):
    
    model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(channels * width * height, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, num_classes)
        )
    
    return model<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1070350592059438, 'fraction_numerical': 0.014859530996052937, 'mean_word_length': 2.691516709511568, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '98989', 'license': 'MIT', 'n_tokens_mistral': 1416, 'n_tokens_neox': 1344, 'n_words': 309, 'path': 'notebooks/mnist_next_steps.ipynb', 'repo_name': 'Data-drone/lightning_vision', 'size': '4307'}