<jupyter_start><jupyter_text>04 MNIST Example with DataLoaders and ConvolutionsIn this notebook we will look at how data loaders can be used in `pycox`.This is particularly useful when working with larger data sets than what is possible to fit in memory, and is an important part of any deep learning framework.As `pycox` is build on [torchtuples](https://github.com/havakv/torchtuples), the same principles applies as for `torchtuples.Model`.For our example, we will consider the [simulation study proposed by Gensheimer and Narasimhan](https://peerj.com/articles/6257/) based on the MNIST data set of handwritten digits. The basic ideas is that each digit represents a survival function, so if we can identify the digit, it is quite straight forward to get good survival estimates.We will use the `LogisticHazard` methods (which [Gensheimer and Narasimhan](https://peerj.com/articles/6257/) refer to as Nnet-survival), with a convolutional network.We will however, consider a slightly different survival function than that of [Gensheimer and Narasimhan](https://peerj.com/articles/6257/), and we will consider all the digits from 0 to 9, while [Gensheimer and Narasimhan](https://peerj.com/articles/6257/) only considered the first 5.<jupyter_code>import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader 

# MNIST is part of torchvision
from torchvision import datasets, transforms


import torchtuples as tt
from pycox.models import LogisticHazard
from pycox.utils import kaplan_meier
from pycox.evaluation import EvalSurv
# for reproducability
np.random.seed(1234)
_ = torch.manual_seed(1234)<jupyter_output><empty_output><jupyter_text>The DatasetWe start by obtaining the MNIST data set with standard preprocessing. The `transform` ensures the data is a `torch.Tensor` and normalize with with a mean and standard deviation.<jupyter_code>transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.1307,), (0.3081,))]
)
mnist_train = datasets.MNIST('.', train=True, download=True,
                             transform=transform)
mnist_test = datasets.MNIST('.', train=False, transform=transform)
_ = plt.imshow(mnist_train[0][0][0].numpy(), cmap='gray')<jupyter_output><empty_output><jupyter_text>SimulationNext we need to simulate the responses corresponding to the images.We draw event times from an exponential distribution with the digit defining the scale parameter$$\beta(\text{digit}) = \frac{365 \cdot \exp(-0.6 \cdot \text{digit})}{\log(1.2)},$$and we censor all times higher than 700.<jupyter_code>def sim_event_times(mnist, max_time=700):
    digits = mnist.targets.numpy()
    betas = 365 * np.exp(-0.6 * digits) / np.log(1.2)
    event_times = np.random.exponential(betas)
    censored = event_times > max_time
    event_times[censored] = max_time
    return tt.tuplefy(event_times, ~censored)<jupyter_output><empty_output><jupyter_text>We simulate a training set and test set, based on the respective MNIST data sets.<jupyter_code>sim_train = sim_event_times(mnist_train)
sim_test = sim_event_times(mnist_test)
sim_train<jupyter_output><empty_output><jupyter_text>VisualizeWe can visualize the survival curves for the 10 digits by applying the Kaplan-Meier estimator to the collection of event times for each digit<jupyter_code>for i in range(10):
    idx = mnist_train.targets.numpy() == i
    kaplan_meier(*sim_train.iloc[idx]).rename(i).plot()
_ = plt.legend()<jupyter_output><empty_output><jupyter_text>Our goal will be to estimate these survival functions from the images.  Label transformsOur simulated event times are drawn in continuous time, so to apply the `LogisticHazard` method, we need to discretize the observations. This can be done with the `label_transform` attribute, and we here use an equidistant grid with 20 grid points.<jupyter_code>labtrans = LogisticHazard.label_transform(20)
target_train = labtrans.fit_transform(*sim_train)
target_test = labtrans.transform(*sim_test)<jupyter_output><empty_output><jupyter_text>The disretization grid is<jupyter_code>labtrans.cuts<jupyter_output><empty_output><jupyter_text>and the discrete targets are<jupyter_code>target_train<jupyter_output><empty_output><jupyter_text>Make DataLoadersTo make a `DataLoader` we first need to create a `Dataset`. The `DataSet` is responsible for the obtaining and transforming the data, while the `DataLoader` is contains a `DataSet` a batch sampler etc.The standard way to create a `Dataset` in PyTorch is by inheriting the `Dataset` class and defining the `__getitem__` method which reads the data for one individual at a time. This also require a `collate_fn` for combining multiple individuals into a batch. The following is an example of this approach, but we will shortly present an alternative approach that is more in line with `torchtuples`.<jupyter_code>class MnistSimDatasetSingle(Dataset):
    """Simulatied data from MNIST. Read a single entry at a time.
    """
    def __init__(self, mnist_dataset, time, event):
        self.mnist_dataset = mnist_dataset
        self.time, self.event = tt.tuplefy(time, event).to_tensor()

    def __len__(self):
        return len(self.mnist_dataset)

    def __getitem__(self, index):
        if type(index) is not int:
            raise ValueError(f"Need `index` to be `int`. Got {type(index)}.")
        img = self.mnist_dataset[index][0]
        return img, (self.time[index], self.event[index])
dataset_train = MnistSimDatasetSingle(mnist_train, *target_train)
dataset_test = MnistSimDatasetSingle(mnist_test, *target_test)
samp = tt.tuplefy(dataset_train[1])
samp.shapes()
samp[1]<jupyter_output><empty_output><jupyter_text>Our dataset gives a nested tuple `(img, (idx_duration, event))`, meaning the default collate in PyTorch does not work. We therefore use `tuplefy` to stack the tensors instead<jupyter_code>def collate_fn(batch):
    """Stacks the entries of a nested tuple"""
    return tt.tuplefy(batch).stack()<jupyter_output><empty_output><jupyter_text>DataLoaderWe can now use the regular pytorch `DataLoader`.Note that you can set the argument `num_workers` in the `DataLoader` to use multiple processes for reading data. Dependent on the system (mac/linux/windows) this can cause some memory issues, so we here use the default `num_workers = 0`.<jupyter_code>batch_size = 128
dl_train = DataLoader(dataset_train, batch_size, shuffle=True, collate_fn=collate_fn)
dl_test = DataLoader(dataset_test, batch_size, shuffle=False, collate_fn=collate_fn)<jupyter_output><empty_output><jupyter_text>If we now investigate a batch, we see that we have the same tuple structure `(img, (idx_durations, events))` but in a batch of size 128.<jupyter_code>batch = next(iter(dl_train))
batch.shapes()
batch.dtypes()<jupyter_output><empty_output><jupyter_text>Dataset with batches (alternative)When working with `torchtuples` it is typically simpler to read a batch at a times. This means that we do not need a `collate_fn`, and all the logic is in the `Dataset`.This approach is not needed, and if you prefer the regular PyTorch `DataLoader`, you can skip this and continue at the Convolutional Network section.<jupyter_code>class MnistSimDatasetBatch(Dataset):
    def __init__(self, mnist_dataset, time, event):
        self.mnist_dataset = mnist_dataset
        self.time, self.event = tt.tuplefy(time, event).to_tensor()

    def __len__(self):
        return len(self.time)

    def __getitem__(self, index):
        if not hasattr(index, '__iter__'):
            index = [index]
        img = [self.mnist_dataset[i][0] for i in index]
        img = torch.stack(img)
        return tt.tuplefy(img, (self.time[index], self.event[index]))
dataset_train = MnistSimDatasetBatch(mnist_train, *target_train)
dataset_test = MnistSimDatasetBatch(mnist_test, *target_test)
samp = dataset_train[[0, 1, 3]]
samp.shapes()<jupyter_output><empty_output><jupyter_text>DataLoaderBatchAs we have a `Dataset` that reads a batch at a time, we cannot use the regular pytorch `DataLoader`.Instead we have to rely on the `DataLoaderBatch` from `torchtuples`, but note that we don't need the `collate_fn`.<jupyter_code>dl_train = tt.data.DataLoaderBatch(dataset_train, batch_size, shuffle=True)
dl_test = tt.data.DataLoaderBatch(dataset_test, batch_size, shuffle=False)
batch = next(iter(dl_train))
batch.shapes()
batch.dtypes()<jupyter_output><empty_output><jupyter_text>We see that the end result is the same as for thte `DataLoader` above, so use the methods you find the simplest.  Convolutional NetworkWe will use a convolutional network with two convolutional layers, global average pooling, and two dense layers. This networks is very basic, so better performance would be expected with a more carefully designed network.<jupyter_code>class Net(nn.Module):
    def __init__(self, out_features):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, 5, 1)
        self.max_pool = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(16, 16, 5, 1)
        self.glob_avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc1 = nn.Linear(16, 16)
        self.fc2 = nn.Linear(16, out_features)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.max_pool(x)
        x = F.relu(self.conv2(x))
        x = self.glob_avg_pool(x)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x
net = Net(labtrans.out_features)
net<jupyter_output><empty_output><jupyter_text>The Logistic-Hazard ModelWe use the `LogisticHazard` with the Adam optimizer with a learning rate of 0.01.<jupyter_code>model = LogisticHazard(net, tt.optim.Adam(0.01), duration_index=labtrans.cuts)<jupyter_output><empty_output><jupyter_text>To verify that the network works as expected we can use the batch from before<jupyter_code>pred = model.predict(batch[0])
pred.shape<jupyter_output><empty_output><jupyter_text>TrainingWe fit the network with `fit_dataloader` and use the `dl_test` to monitor the test performance. It should go without saying that, in practice, we need a validation set separate from the test set when we use early stopping, but this is just an illustrative example.<jupyter_code>callbacks = [tt.cb.EarlyStopping(patience=5)]
epochs = 50
verbose = True
log = model.fit_dataloader(dl_train, epochs, callbacks, verbose, val_dataloader=dl_test)
_ = log.plot()<jupyter_output><empty_output><jupyter_text>PredictionTo predict, we need a data loader that only gives the images and not the targets. We therefore need to create a new `Dataset` for this purpose.<jupyter_code>class MnistSimInput(Dataset):
    def __init__(self, mnist_dataset):
        self.mnist_dataset = mnist_dataset

    def __len__(self):
        return len(self.mnist_dataset)

    def __getitem__(self, index):
        img = self.mnist_dataset[index][0]
        return img
dataset_test_x = MnistSimInput(mnist_test)
dl_test_x = DataLoader(dataset_test_x, batch_size, shuffle=False)
next(iter(dl_test_x)).shape<jupyter_output><empty_output><jupyter_text>(alternative)Alternatively, if you have used the batch method, we can use the method `dataloader_input_only` to create this `Dataloader` from `dl_test`.<jupyter_code>dl_test_x = tt.data.dataloader_input_only(dl_test)
next(iter(dl_test_x)).shape<jupyter_output><empty_output><jupyter_text>Survial predictionsWe can obtain survival prediction in the regular manner, and one can include the `interpolation` if wanted.<jupyter_code>surv = model.predict_surv_df(dl_test_x)<jupyter_output><empty_output><jupyter_text>ResultsWe compute the average survival predictions for each digit in the test set<jupyter_code>for i in range(10):
    idx = mnist_test.targets.numpy() == i
    surv.loc[:, idx].mean(axis=1).rename(i).plot()
_ = plt.legend()<jupyter_output><empty_output><jupyter_text>and find that they are quite similar to the Kaplan-Meier estimates!<jupyter_code>for i in range(10):
    idx = mnist_test.targets.numpy() == i
    kaplan_meier(*sim_test.iloc[idx]).rename(i).plot()
_ = plt.legend()<jupyter_output><empty_output><jupyter_text>Concordance and Brier score<jupyter_code>surv = model.interpolate(10).predict_surv_df(dl_test_x)
ev = EvalSurv(surv, *sim_test, 'km')
ev.concordance_td()
time_grid = np.linspace(0, sim_test[0].max())
ev.integrated_brier_score(time_grid)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.094580405459979, 'fraction_numerical': 0.011065342056376706, 'mean_word_length': 5.359527478171546, 'pattern_counts': {'":': 0, '<': 97, '<?xml version=': 0, '>': 98, 'https://': 5, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '24', 'dir': 'jupyter-structured-clean-dedup', 'id': '484773', 'license': 'BSD-2-Clause', 'n_tokens_mistral': 4072, 'n_tokens_neox': 3922, 'n_words': 1185, 'path': 'examples/04_mnist_dataloaders_cnn.ipynb', 'repo_name': 'BardiaKh/pycox', 'size': '12381'}