import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
import torch
from tqdm import tqdm

import pytorch_lightning as pl

ToC

1. Data
   1. Load Data
   2. Dataset
   3. Dataloader

2. Model

   1. Vanilla pytorch model
   2. Lightning model
   3. Integrated training loop
   4. Log training loss and accuracy
   5. Integrate training dataloder into model
   6. LR Find
   7. LR Schedulars & log LR

3. Visualize

1. Data

Load data

Data used for this tutorial is "test" set of FashionMNIST, downloadable from https://github.com/zalandoresearch/fashion-mnist

• Images http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz
• Labels http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz

label_list = [
    'T-shirt/top',
    'Trouser',
    'Pullover',
    'Dress',
    'Coat',
    'Sandal',
    'Shirt',
    'Sneaker',
    'Bag',
    'Ankle boot'
    ]

labels = np.frombuffer(open('./t10k-labels-idx1-ubyte', 'rb').read(), dtype=np.uint8,
                               offset=8)

images = np.frombuffer(open('./t10k-images-idx3-ubyte', 'rb').read(), dtype=np.uint8,
                               offset=16).reshape(len(labels), 784)

images.shape, labels.shape

((10000, 784), (10000,))

idx = np.random.randint(10000)
plt.imshow(images[idx].reshape(28, 28), cmap='gray')
label_list[labels[idx]]

Dataset

class FashionDS(Dataset):
    index = None
    def __init__(self, train = True):
        self.imgs = images
        self.labels = labels
        self.label_lst = label_list
        self.train = train
        self.train_fraction = len(labels) * 8 // 10
        if FashionDS.index is None:
            FashionDS.index = np.random.permutation(np.arange(len(self.labels)))
        
    def __len__(self):
        if self.train:
            return self.train_fraction
        else:
            return len(self.labels) - self.train_fraction
    
    def __getitem__(self, idx):
        if self.train:
            idx2 = FashionDS.index[idx]
        else:
            idx2 = FashionDS.index[self.fraction + idx]
        return self.imgs[idx2].reshape(28, 28), self.labels[idx2]

ds = FashionDS()
print(len(ds))
idx = np.random.randint(len(ds))
x, y = ds[idx]
plt.imshow(x, cmap='gray')
ds.label_lst[y]

8000

'Pullover'

valds = FashionDS(False)
print(len(valds))
idx = np.random.randint(len(valds))
x, y = ds[idx]
plt.imshow(x, cmap='gray')
ds.label_lst[y]

2000

'Bag'

Dataloader

dl = DataLoader(ds, batch_size=16, shuffle=True)

x, y = next(iter(dl))
x.shape, y.shape

(torch.Size([16, 28, 28]), torch.Size([16]))

2. Model

Vanilla pytorch model

class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 5, 3)
        self.linear = torch.nn.Linear(26*26*5, 10)
        
    def forward(self, x):
        out = self.conv(x)
        out = torch.nn.functional.relu(out)
        out = out.view(-1, 26*26*5)
        out = self.linear(out)
        return out

net = Net()

x, y = next(iter(dl))
x = x.float() / 255.
yhat = net(x[:, None, ...])
yhat.shape

torch.Size([16, 10])

Loss

torch.nn.functional.cross_entropy(yhat, y.long()).item()

2.314405679702759

optimizer

opt = torch.optim.SGD(net.parameters(), 3e-4, 0.9)

Train

loss_lst = []
acc = []
for epoch in range(5):
    print(epoch)
    for x, y in tqdm(dl):
        
        opt.zero_grad()
        x = x.float() / 255.
        yhat = net(x[:, None, ...])
        acc += [(yhat.argmax(dim=1) == y).sum()*1.0 / len(y)]
            
        loss = torch.nn.functional.cross_entropy(yhat, y.long())
        loss_lst += [loss.item()]
        loss.backward()
        
        opt.step()        

plt.plot(loss_lst)

plt.plot(acc)

np.mean(acc)

0.81164

Lightning model

class Net(pl.LightningModule):                    ## changed
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 5, 3)
        self.linear = torch.nn.Linear(26*26*5, 10)
        
    def forward(self, x):
        out = self.conv(x)
        out = torch.nn.functional.relu(out)
        out = out.view(-1, 26*26*5)
        out = self.linear(out)
        return out

net = Net()

Sanity check

x, y = next(iter(dl))
x = x.float() / 255.
yhat = net(x[:, None, ...])
yhat.shape

torch.nn.functional.cross_entropy(yhat, y.long()).item()

2.3112006187438965

opt = torch.optim.SGD(net.parameters(), 3e-4, 0.9)

loss_lst = []
acc = []
for epoch in range(5):
    print(epoch)
    for x, y in tqdm(dl):
        
        opt.zero_grad()
        x = x.float() / 255.
        yhat = net(x[:, None, ...])
        acc += [(yhat.argmax(dim=1) == y).sum()*1.0 / len(y)]
            
        loss = torch.nn.functional.cross_entropy(yhat, y.long())
        loss_lst += [loss.item()]
        loss.backward()
        
        opt.step()        

plt.plot(loss_lst)

plt.plot(acc)

np.mean(acc)

0.6867

Integrated training loop

class Net(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 5, 3)
        self.linear = torch.nn.Linear(26*26*5, 10)
        
    def forward(self, x):
        out = self.conv(x)
        out = torch.nn.functional.relu(out)
        out = out.view(-1, 26*26*5)
        out = self.linear(out)
        return out
    
    def training_step(self, batch, batch_id):                   ## changed
        x, y = batch
        x = x.float() / 255.
        yhat = self(x[:, None, ...])
        loss = torch.nn.functional.cross_entropy(yhat, y.long())
        return loss
    
    def configure_optimizers(self):                             ## changed
        opt = torch.optim.SGD(net.parameters(), 3e-4, 0.9)      
        return opt

net = Net()

Sanity check

x, y = next(iter(dl))
x = x.float() / 255.
yhat = net(x[:, None, ...])
yhat.shape

torch.nn.functional.cross_entropy(yhat, y.long()).item()

2.3086438179016113

trainer = pl.Trainer()            ## changed
trainer.max_epochs = 5            ## changed
trainer.fit(net, dl)              ## changed

Log training loss and accuracy

class Net(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 5, 3)
        self.linear = torch.nn.Linear(26*26*5, 10)
        
    def forward(self, x):
        out = self.conv(x)
        out = torch.nn.functional.relu(out)
        out = out.view(-1, 26*26*5)
        out = self.linear(out)
        return out
    
    def training_step(self, batch, batch_id):
        x, y = batch
        x = x.float() / 255.
        yhat = self(x[:, None, ...])
        acc = (yhat.argmax(dim=1) == y).sum()*1.0 / len(y)
        loss = torch.nn.functional.cross_entropy(yhat, y.long())
        self.log('train loss', loss.item())                     ## changed
        self.log('train acc', acc)                              ## changed
        return loss 
    
    def configure_optimizers(self):
        opt = torch.optim.SGD(net.parameters(), 3e-4, 0.9)
        return opt

net = Net()

Sanity check

x, y = next(iter(dl))
x = x.float() / 255.
yhat = net(x[:, None, ...])
yhat.shape

torch.nn.functional.cross_entropy(yhat, y.long()).item()

2.2562389373779297

trainer = pl.Trainer()
trainer.max_epochs = 5
trainer.fit(net, dl)

Integrate training dataloder into model

class Net(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 5, 3)
        self.linear = torch.nn.Linear(26*26*5, 10)
        
    def forward(self, x):
        out = self.conv(x)
        out = torch.nn.functional.relu(out)
        out = out.view(-1, 26*26*5)
        out = self.linear(out)
        return out
    
    def training_step(self, batch, batch_id):
        x, y = batch
        x = x.float() / 255.
        yhat = self(x[:, None, ...])
        acc = (yhat.argmax(dim=1) == y).sum()*1.0 / len(y)
        loss = torch.nn.functional.cross_entropy(yhat, y.long())
        self.log('train loss', loss.item())
        self.log('train acc', acc)
        return loss
    
    def configure_optimizers(self):
        opt = torch.optim.SGD(net.parameters(), 3e-4, 0.9)
        return opt
        
    def train_dataloader(self):                       ## changed
        return DataLoader(FashionDS(), batch_size=64)

net = Net()

Sanity check

x, y = next(iter(net.train_dataloader()))
x = x.float() / 255.
yhat = net(x[:, None, ...])
yhat.shape

torch.nn.functional.cross_entropy(yhat, y.long()).item()

2.28153133392334

trainer = pl.Trainer()
trainer.max_epochs = 5
trainer.fit(net)        ## changed

LR Find

class Net(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 5, 3)
        self.linear = torch.nn.Linear(26*26*5, 10)
        
    def forward(self, x):
        out = self.conv(x)
        out = torch.nn.functional.relu(out)
        out = out.view(-1, 26*26*5)
        out = self.linear(out)
        return out
    
    def training_step(self, batch, batch_id):
        x, y = batch
        x = x.float() / 255.
        yhat = self(x[:, None, ...])
        acc = (yhat.argmax(dim=1) == y).sum()*1.0 / len(y)
        loss = torch.nn.functional.cross_entropy(yhat, y.long())
        self.log('train loss', loss.item())
        self.log('train acc', acc)
        return loss
    
    def configure_optimizers(self):
        self.lr = 3e-4
        opt = torch.optim.SGD(net.parameters(), self.lr, 0.9)       ## changed
        return opt
        
    def train_dataloader(self):
        return DataLoader(FashionDS(), batch_size=64)

net = Net()

Sanity check

x, y = next(iter(net.train_dataloader()))
x = x.float() / 255.
yhat = net(x[:, None, ...])
yhat.shape

torch.nn.functional.cross_entropy(yhat, y.long()).item()

2.321258306503296

trainer = pl.Trainer()
trainer.max_epochs = 10

lrf = trainer.tuner.lr_find(net)  ## changed
fig = lrf.plot(suggest=True)      ## changed
fig.show()                        ## changed

net.hparams.lr = 1e-2            ## changed

trainer.fit(net)

LR Schedulars & log LR

class Net(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(1, 5, 3)
        self.linear = torch.nn.Linear(26*26*5, 10)
        
    def forward(self, x):
        out = self.conv(x)
        out = torch.nn.functional.relu(out)
        out = out.view(-1, 26*26*5)
        out = self.linear(out)
        return out
    
    def training_step(self, batch, batch_id):
        x, y = batch
        x = x.float() / 255.
        yhat = self(x[:, None, ...])
        acc = (yhat.argmax(dim=1) == y).sum()*1.0 / len(y)
        loss = torch.nn.functional.cross_entropy(yhat, y.long())
        self.log('train loss', loss.item())
        self.log('train acc', acc)
        return loss
    
    def configure_optimizers(self):
        self.lr = 3e-4
        opt = torch.optim.SGD(net.parameters(), self.lr, 0.9)
        sched = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, 200) ## changed
        return [opt], [{                                                       ## changed
                'scheduler': sched,                                            ## changed
                'interval': 'step', # The unit of the scheduler's step size    ## changed
            }]
        
    def train_dataloader(self):
        return DataLoader(FashionDS(), batch_size=64)

net = Net()

Sanity check

x, y = next(iter(net.train_dataloader()))
x = x.float() / 255.
yhat = net(x[:, None, ...])
yhat.shape

torch.nn.functional.cross_entropy(yhat, y.long()).item()

2.314138889312744

trainer = pl.Trainer(callbacks=[pl.callbacks.LearningRateMonitor('step')])  ## changed
trainer.max_epochs = 3
net.hparams.lr = 1e-2
trainer.fit(net)

3. Visualize

%load_ext tensorboard
%tensorboard --logdir ./lightning_logs