PyTorch 2.2 中文官方教程（五）

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt 

epsilons = [0, .05, .1, .15, .2, .25, .3]
pretrained_model = "data/lenet_mnist_model.pth"
use_cuda=True
# Set random seed for reproducibility
torch.manual_seed(42) 

<torch._C.Generator object at 0x7f6b149d3070> 

# LeNet Model definition
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

# MNIST Test dataset and dataloader declaration
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=False, download=True, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,)),
            ])),
        batch_size=1, shuffle=True)

# Define what device we are using
print("CUDA Available: ",torch.cuda.is_available())
device = torch.device("cuda" if use_cuda and torch.cuda.is_available() else "cpu")

# Initialize the network
model = Net().to(device)

# Load the pretrained model
model.load_state_dict(torch.load(pretrained_model, map_location=device))

# Set the model in evaluation mode. In this case this is for the Dropout layers
model.eval() 

Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ../data/MNIST/raw/train-images-idx3-ubyte.gz

  0%|          | 0/9912422 [00:00<?, ?it/s]
100%|##########| 9912422/9912422 [00:00<00:00, 436275131.90it/s]
Extracting ../data/MNIST/raw/train-images-idx3-ubyte.gz to ../data/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ../data/MNIST/raw/train-labels-idx1-ubyte.gz

  0%|          | 0/28881 [00:00<?, ?it/s]
100%|##########| 28881/28881 [00:00<00:00, 35440518.97it/s]
Extracting ../data/MNIST/raw/train-labels-idx1-ubyte.gz to ../data/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ../data/MNIST/raw/t10k-images-idx3-ubyte.gz

  0%|          | 0/1648877 [00:00<?, ?it/s]
100%|##########| 1648877/1648877 [00:00<00:00, 251450385.28it/s]
Extracting ../data/MNIST/raw/t10k-images-idx3-ubyte.gz to ../data/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ../data/MNIST/raw/t10k-labels-idx1-ubyte.gz

  0%|          | 0/4542 [00:00<?, ?it/s]
100%|##########| 4542/4542 [00:00<00:00, 36286721.46it/s]
Extracting ../data/MNIST/raw/t10k-labels-idx1-ubyte.gz to ../data/MNIST/raw

CUDA Available:  True

Net(
  (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
  (dropout1): Dropout(p=0.25, inplace=False)
  (dropout2): Dropout(p=0.5, inplace=False)
  (fc1): Linear(in_features=9216, out_features=128, bias=True)
  (fc2): Linear(in_features=128, out_features=10, bias=True)
) 

# FGSM attack code
def fgsm_attack(image, epsilon, data_grad):
    # Collect the element-wise sign of the data gradient
    sign_data_grad = data_grad.sign()
    # Create the perturbed image by adjusting each pixel of the input image
    perturbed_image = image + epsilon*sign_data_grad
    # Adding clipping to maintain [0,1] range
    perturbed_image = torch.clamp(perturbed_image, 0, 1)
    # Return the perturbed image
    return perturbed_image

# restores the tensors to their original scale
def denorm(batch, mean=[0.1307], std=[0.3081]):
  """
 Convert a batch of tensors to their original scale.

 Args:
 batch (torch.Tensor): Batch of normalized tensors.
 mean (torch.Tensor or list): Mean used for normalization.
 std (torch.Tensor or list): Standard deviation used for normalization.

 Returns:
 torch.Tensor: batch of tensors without normalization applied to them.
 """
    if isinstance(mean, list):
        mean = torch.tensor(mean).to(device)
    if isinstance(std, list):
        std = torch.tensor(std).to(device)

    return batch * std.view(1, -1, 1, 1) + mean.view(1, -1, 1, 1) 

def test( model, device, test_loader, epsilon ):

    # Accuracy counter
    correct = 0
    adv_examples = []

    # Loop over all examples in test set
    for data, target in test_loader:

        # Send the data and label to the device
        data, target = data.to(device), target.to(device)

        # Set requires_grad attribute of tensor. Important for Attack
        data.requires_grad = True

        # Forward pass the data through the model
        output = model(data)
        init_pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability

        # If the initial prediction is wrong, don't bother attacking, just move on
        if init_pred.item() != target.item():
            continue

        # Calculate the loss
        loss = F.nll_loss(output, target)

        # Zero all existing gradients
        model.zero_grad()

        # Calculate gradients of model in backward pass
        loss.backward()

        # Collect ``datagrad``
        data_grad = data.grad.data

        # Restore the data to its original scale
        data_denorm = denorm(data)

        # Call FGSM Attack
        perturbed_data = fgsm_attack(data_denorm, epsilon, data_grad)

        # Reapply normalization
        perturbed_data_normalized = transforms.Normalize((0.1307,), (0.3081,))(perturbed_data)

        # Re-classify the perturbed image
        output = model(perturbed_data_normalized)

        # Check for success
        final_pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
        if final_pred.item() == target.item():
            correct += 1
            # Special case for saving 0 epsilon examples
            if epsilon == 0 and len(adv_examples) < 5:
                adv_ex = perturbed_data.squeeze().detach().cpu().numpy()
                adv_examples.append( (init_pred.item(), final_pred.item(), adv_ex) )
        else:
            # Save some adv examples for visualization later
            if len(adv_examples) < 5:
                adv_ex = perturbed_data.squeeze().detach().cpu().numpy()
                adv_examples.append( (init_pred.item(), final_pred.item(), adv_ex) )

    # Calculate final accuracy for this epsilon
    final_acc = correct/float(len(test_loader))
    print(f"Epsilon: {epsilon}\tTest Accuracy = {correct} / {len(test_loader)} = {final_acc}")

    # Return the accuracy and an adversarial example
    return final_acc, adv_examples 

accuracies = []
examples = []

# Run test for each epsilon
for eps in epsilons:
    acc, ex = test(model, device, test_loader, eps)
    accuracies.append(acc)
    examples.append(ex) 

Epsilon: 0      Test Accuracy = 9912 / 10000 = 0.9912
Epsilon: 0.05   Test Accuracy = 9605 / 10000 = 0.9605
Epsilon: 0.1    Test Accuracy = 8743 / 10000 = 0.8743
Epsilon: 0.15   Test Accuracy = 7111 / 10000 = 0.7111
Epsilon: 0.2    Test Accuracy = 4877 / 10000 = 0.4877
Epsilon: 0.25   Test Accuracy = 2717 / 10000 = 0.2717
Epsilon: 0.3    Test Accuracy = 1418 / 10000 = 0.1418 

plt.figure(figsize=(5,5))
plt.plot(epsilons, accuracies, "*-")
plt.yticks(np.arange(0, 1.1, step=0.1))
plt.xticks(np.arange(0, .35, step=0.05))
plt.title("Accuracy vs Epsilon")
plt.xlabel("Epsilon")
plt.ylabel("Accuracy")
plt.show() 

# Plot several examples of adversarial samples at each epsilon
cnt = 0
plt.figure(figsize=(8,10))
for i in range(len(epsilons)):
    for j in range(len(examples[i])):
        cnt += 1
        plt.subplot(len(epsilons),len(examples[0]),cnt)
        plt.xticks([], [])
        plt.yticks([], [])
        if j == 0:
            plt.ylabel(f"Eps: {epsilons[i]}", fontsize=14)
        orig,adv,ex = examples[i][j]
        plt.title(f"{orig} -> {adv}")
        plt.imshow(ex, cmap="gray")
plt.tight_layout()
plt.show() 

#%matplotlib inline
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML

# Set random seed for reproducibility
manualSeed = 999
#manualSeed = random.randint(1, 10000) # use if you want new results
print("Random Seed: ", manualSeed)
random.seed(manualSeed)
torch.manual_seed(manualSeed)
torch.use_deterministic_algorithms(True) # Needed for reproducible results 

Random Seed:  999 

# Root directory for dataset
dataroot = "data/celeba"

# Number of workers for dataloader
workers = 2

# Batch size during training
batch_size = 128

# Spatial size of training images. All images will be resized to this
#   size using a transformer.
image_size = 64

# Number of channels in the training images. For color images this is 3
nc = 3

# Size of z latent vector (i.e. size of generator input)
nz = 100

# Size of feature maps in generator
ngf = 64

# Size of feature maps in discriminator
ndf = 64

# Number of training epochs
num_epochs = 5

# Learning rate for optimizers
lr = 0.0002

# Beta1 hyperparameter for Adam optimizers
beta1 = 0.5

# Number of GPUs available. Use 0 for CPU mode.
ngpu = 1 

/path/to/celeba
  ->  img_align_celeba
  ->  188242.jpg
  ->  173822.jpg
  ->  284702.jpg
  ->  537394.jpg
  ... 

# We can use an image folder dataset the way we have it setup.
# Create the dataset
dataset = dset.ImageFolder(root=dataroot,
                           transform=transforms.Compose([
                               transforms.Resize(image_size),
                               transforms.CenterCrop(image_size),
                               transforms.ToTensor(),
                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                           ]))
# Create the dataloader
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
                                         shuffle=True, num_workers=workers)

# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")

# Plot some training images
real_batch = next(iter(dataloader))
plt.figure(figsize=(8,8))
plt.axis("off")
plt.title("Training Images")
plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
plt.show() 

# custom weights initialization called on ``netG`` and ``netD``
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0) 

# Generator Code

class Generator(nn.Module):
    def __init__(self, ngpu):
        super(Generator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. ``(ngf*8) x 4 x 4``
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. ``(ngf*4) x 8 x 8``
            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. ``(ngf*2) x 16 x 16``
            nn.ConvTranspose2d( ngf * 2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. ``(ngf) x 32 x 32``
            nn.ConvTranspose2d( ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. ``(nc) x 64 x 64``
        )

    def forward(self, input):
        return self.main(input) 

# Create the generator
netG = Generator(ngpu).to(device)

# Handle multi-GPU if desired
if (device.type == 'cuda') and (ngpu > 1):
    netG = nn.DataParallel(netG, list(range(ngpu)))

# Apply the ``weights_init`` function to randomly initialize all weights
#  to ``mean=0``, ``stdev=0.02``.
netG.apply(weights_init)

# Print the model
print(netG) 

Generator(
  (main): Sequential(
    (0): ConvTranspose2d(100, 512, kernel_size=(4, 4), stride=(1, 1), bias=False)
    (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
    (3): ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (5): ReLU(inplace=True)
    (6): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (7): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (8): ReLU(inplace=True)
    (9): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (10): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (11): ReLU(inplace=True)
    (12): ConvTranspose2d(64, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (13): Tanh()
  )
) 

class Discriminator(nn.Module):
    def __init__(self, ngpu):
        super(Discriminator, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is ``(nc) x 64 x 64``
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. ``(ndf) x 32 x 32``
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. ``(ndf*2) x 16 x 16``
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. ``(ndf*4) x 8 x 8``
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. ``(ndf*8) x 4 x 4``
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input) 

# Create the Discriminator
netD = Discriminator(ngpu).to(device)

# Handle multi-GPU if desired
if (device.type == 'cuda') and (ngpu > 1):
    netD = nn.DataParallel(netD, list(range(ngpu)))

# Apply the ``weights_init`` function to randomly initialize all weights
# like this: ``to mean=0, stdev=0.2``.
netD.apply(weights_init)

# Print the model
print(netD) 

Discriminator(
  (main): Sequential(
    (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (1): LeakyReLU(negative_slope=0.2, inplace=True)
    (2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (4): LeakyReLU(negative_slope=0.2, inplace=True)
    (5): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (6): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (7): LeakyReLU(negative_slope=0.2, inplace=True)
    (8): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (9): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (10): LeakyReLU(negative_slope=0.2, inplace=True)
    (11): Conv2d(512, 1, kernel_size=(4, 4), stride=(1, 1), bias=False)
    (12): Sigmoid()
  )
) 

# Initialize the ``BCELoss`` function
criterion = nn.BCELoss()

# Create batch of latent vectors that we will use to visualize
#  the progression of the generator
fixed_noise = torch.randn(64, nz, 1, 1, device=device)

# Establish convention for real and fake labels during training
real_label = 1.
fake_label = 0.

# Setup Adam optimizers for both G and D
optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999)) 

# Training Loop

# Lists to keep track of progress
img_list = []
G_losses = []
D_losses = []
iters = 0

print("Starting Training Loop...")
# For each epoch
for epoch in range(num_epochs):
    # For each batch in the dataloader
    for i, data in enumerate(dataloader, 0):

        ############################
        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
        ###########################
        ## Train with all-real batch
        netD.zero_grad()
        # Format batch
        real_cpu = data[0].to(device)
        b_size = real_cpu.size(0)
        label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
        # Forward pass real batch through D
        output = netD(real_cpu).view(-1)
        # Calculate loss on all-real batch
        errD_real = criterion(output, label)
        # Calculate gradients for D in backward pass
        errD_real.backward()
        D_x = output.mean().item()

        ## Train with all-fake batch
        # Generate batch of latent vectors
        noise = torch.randn(b_size, nz, 1, 1, device=device)
        # Generate fake image batch with G
        fake = netG(noise)
        label.fill_(fake_label)
        # Classify all fake batch with D
        output = netD(fake.detach()).view(-1)
        # Calculate D's loss on the all-fake batch
        errD_fake = criterion(output, label)
        # Calculate the gradients for this batch, accumulated (summed) with previous gradients
        errD_fake.backward()
        D_G_z1 = output.mean().item()
        # Compute error of D as sum over the fake and the real batches
        errD = errD_real + errD_fake
        # Update D
        optimizerD.step()

        ############################
        # (2) Update G network: maximize log(D(G(z)))
        ###########################
        netG.zero_grad()
        label.fill_(real_label)  # fake labels are real for generator cost
        # Since we just updated D, perform another forward pass of all-fake batch through D
        output = netD(fake).view(-1)
        # Calculate G's loss based on this output
        errG = criterion(output, label)
        # Calculate gradients for G
        errG.backward()
        D_G_z2 = output.mean().item()
        # Update G
        optimizerG.step()

        # Output training stats
        if i % 50 == 0:
            print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
                  % (epoch, num_epochs, i, len(dataloader),
                     errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))

        # Save Losses for plotting later
        G_losses.append(errG.item())
        D_losses.append(errD.item())

        # Check how the generator is doing by saving G's output on fixed_noise
        if (iters % 500 == 0) or ((epoch == num_epochs-1) and (i == len(dataloader)-1)):
            with torch.no_grad():
                fake = netG(fixed_noise).detach().cpu()
            img_list.append(vutils.make_grid(fake, padding=2, normalize=True))

        iters += 1 

Starting Training Loop...
[0/5][0/1583]   Loss_D: 1.4640  Loss_G: 6.9360  D(x): 0.7143    D(G(z)): 0.5877 / 0.0017
[0/5][50/1583]  Loss_D: 0.0174  Loss_G: 23.7368 D(x): 0.9881    D(G(z)): 0.0000 / 0.0000
[0/5][100/1583] Loss_D: 0.5983  Loss_G: 9.9471  D(x): 0.9715    D(G(z)): 0.3122 / 0.0003
[0/5][150/1583] Loss_D: 0.4940  Loss_G: 5.6772  D(x): 0.7028    D(G(z)): 0.0241 / 0.0091
[0/5][200/1583] Loss_D: 0.5931  Loss_G: 7.1186  D(x): 0.9423    D(G(z)): 0.3016 / 0.0018
[0/5][250/1583] Loss_D: 0.3846  Loss_G: 3.2697  D(x): 0.7663    D(G(z)): 0.0573 / 0.0739
[0/5][300/1583] Loss_D: 1.3306  Loss_G: 8.3204  D(x): 0.8768    D(G(z)): 0.6353 / 0.0009
[0/5][350/1583] Loss_D: 0.6451  Loss_G: 6.0499  D(x): 0.9025    D(G(z)): 0.3673 / 0.0060
[0/5][400/1583] Loss_D: 0.4211  Loss_G: 3.7316  D(x): 0.8407    D(G(z)): 0.1586 / 0.0392
[0/5][450/1583] Loss_D: 0.6569  Loss_G: 2.4818  D(x): 0.6437    D(G(z)): 0.0858 / 0.1129
[0/5][500/1583] Loss_D: 1.2208  Loss_G: 2.9943  D(x): 0.4179    D(G(z)): 0.0109 / 0.1133
[0/5][550/1583] Loss_D: 0.3400  Loss_G: 4.7669  D(x): 0.9135    D(G(z)): 0.1922 / 0.0145
[0/5][600/1583] Loss_D: 0.5756  Loss_G: 4.8500  D(x): 0.9189    D(G(z)): 0.3193 / 0.0187
[0/5][650/1583] Loss_D: 0.2470  Loss_G: 4.1606  D(x): 0.9460    D(G(z)): 0.1545 / 0.0250
[0/5][700/1583] Loss_D: 0.3887  Loss_G: 4.1884  D(x): 0.8518    D(G(z)): 0.1562 / 0.0297
[0/5][750/1583] Loss_D: 0.5353  Loss_G: 4.1742  D(x): 0.8034    D(G(z)): 0.1958 / 0.0302
[0/5][800/1583] Loss_D: 0.3213  Loss_G: 5.8919  D(x): 0.9076    D(G(z)): 0.1572 / 0.0065
[0/5][850/1583] Loss_D: 0.8850  Loss_G: 7.4333  D(x): 0.9258    D(G(z)): 0.4449 / 0.0017
[0/5][900/1583] Loss_D: 1.2624  Loss_G: 10.0392 D(x): 0.9896    D(G(z)): 0.6361 / 0.0002
[0/5][950/1583] Loss_D: 0.8802  Loss_G: 6.9221  D(x): 0.5527    D(G(z)): 0.0039 / 0.0045
[0/5][1000/1583]        Loss_D: 0.5799  Loss_G: 3.1800  D(x): 0.7062    D(G(z)): 0.0762 / 0.0884
[0/5][1050/1583]        Loss_D: 0.9647  Loss_G: 6.6894  D(x): 0.9429    D(G(z)): 0.5270 / 0.0035
[0/5][1100/1583]        Loss_D: 0.5624  Loss_G: 3.6715  D(x): 0.7944    D(G(z)): 0.2069 / 0.0445
[0/5][1150/1583]        Loss_D: 0.6205  Loss_G: 4.8995  D(x): 0.8634    D(G(z)): 0.3046 / 0.0169
[0/5][1200/1583]        Loss_D: 0.2569  Loss_G: 4.2945  D(x): 0.9455    D(G(z)): 0.1528 / 0.0255
[0/5][1250/1583]        Loss_D: 0.4921  Loss_G: 3.2500  D(x): 0.8152    D(G(z)): 0.1892 / 0.0753
[0/5][1300/1583]        Loss_D: 0.4068  Loss_G: 3.7702  D(x): 0.8153    D(G(z)): 0.1335 / 0.0472
[0/5][1350/1583]        Loss_D: 1.1704  Loss_G: 7.3408  D(x): 0.9443    D(G(z)): 0.5863 / 0.0022
[0/5][1400/1583]        Loss_D: 0.6111  Loss_G: 2.2676  D(x): 0.6714    D(G(z)): 0.0793 / 0.1510
[0/5][1450/1583]        Loss_D: 0.7817  Loss_G: 4.0744  D(x): 0.7915    D(G(z)): 0.3573 / 0.0242
[0/5][1500/1583]        Loss_D: 0.7177  Loss_G: 1.9253  D(x): 0.5770    D(G(z)): 0.0257 / 0.1909
[0/5][1550/1583]        Loss_D: 0.4518  Loss_G: 2.8314  D(x): 0.7991    D(G(z)): 0.1479 / 0.0885
[1/5][0/1583]   Loss_D: 0.4267  Loss_G: 4.5150  D(x): 0.8976    D(G(z)): 0.2401 / 0.0196
[1/5][50/1583]  Loss_D: 0.5106  Loss_G: 2.7800  D(x): 0.7073    D(G(z)): 0.0663 / 0.0932
[1/5][100/1583] Loss_D: 0.6300  Loss_G: 1.8648  D(x): 0.6557    D(G(z)): 0.0756 / 0.2118
[1/5][150/1583] Loss_D: 1.1727  Loss_G: 5.1536  D(x): 0.8397    D(G(z)): 0.5261 / 0.0125
[1/5][200/1583] Loss_D: 0.4675  Loss_G: 2.9615  D(x): 0.7645    D(G(z)): 0.1400 / 0.0780
[1/5][250/1583] Loss_D: 0.7938  Loss_G: 3.1614  D(x): 0.6958    D(G(z)): 0.2248 / 0.0678
[1/5][300/1583] Loss_D: 0.9869  Loss_G: 5.9243  D(x): 0.9619    D(G(z)): 0.5349 / 0.0063
[1/5][350/1583] Loss_D: 0.5178  Loss_G: 3.0236  D(x): 0.7795    D(G(z)): 0.1769 / 0.0700
[1/5][400/1583] Loss_D: 1.4509  Loss_G: 2.7187  D(x): 0.3278    D(G(z)): 0.0133 / 0.1273
[1/5][450/1583] Loss_D: 0.5530  Loss_G: 4.8110  D(x): 0.9151    D(G(z)): 0.3237 / 0.0160
[1/5][500/1583] Loss_D: 0.4621  Loss_G: 4.1158  D(x): 0.8720    D(G(z)): 0.2278 / 0.0293
[1/5][550/1583] Loss_D: 0.4987  Loss_G: 4.0199  D(x): 0.8533    D(G(z)): 0.2367 / 0.0287
[1/5][600/1583] Loss_D: 1.0630  Loss_G: 4.6502  D(x): 0.9145    D(G(z)): 0.5018 / 0.0218
[1/5][650/1583] Loss_D: 0.6081  Loss_G: 4.3172  D(x): 0.8670    D(G(z)): 0.3312 / 0.0221
[1/5][700/1583] Loss_D: 0.4703  Loss_G: 2.4900  D(x): 0.7538    D(G(z)): 0.1245 / 0.1188
[1/5][750/1583] Loss_D: 0.4827  Loss_G: 2.2941  D(x): 0.7372    D(G(z)): 0.1105 / 0.1300
[1/5][800/1583] Loss_D: 0.4013  Loss_G: 3.8850  D(x): 0.8895    D(G(z)): 0.2179 / 0.0324
[1/5][850/1583] Loss_D: 0.7245  Loss_G: 1.9088  D(x): 0.6100    D(G(z)): 0.0950 / 0.1898
[1/5][900/1583] Loss_D: 0.8372  Loss_G: 1.2346  D(x): 0.5232    D(G(z)): 0.0332 / 0.3633
[1/5][950/1583] Loss_D: 0.5561  Loss_G: 3.2048  D(x): 0.7660    D(G(z)): 0.2035 / 0.0594
[1/5][1000/1583]        Loss_D: 0.6859  Loss_G: 1.6347  D(x): 0.5764    D(G(z)): 0.0435 / 0.2540
[1/5][1050/1583]        Loss_D: 0.6785  Loss_G: 4.3244  D(x): 0.9066    D(G(z)): 0.3835 / 0.0203
[1/5][1100/1583]        Loss_D: 0.4835  Loss_G: 2.4080  D(x): 0.7428    D(G(z)): 0.1073 / 0.1147
[1/5][1150/1583]        Loss_D: 0.5507  Loss_G: 2.5400  D(x): 0.7857    D(G(z)): 0.2182 / 0.1092
[1/5][1200/1583]        Loss_D: 0.6054  Loss_G: 3.4802  D(x): 0.8263    D(G(z)): 0.2934 / 0.0441
[1/5][1250/1583]        Loss_D: 0.4788  Loss_G: 2.3533  D(x): 0.7872    D(G(z)): 0.1698 / 0.1327
[1/5][1300/1583]        Loss_D: 0.5314  Loss_G: 2.7018  D(x): 0.8273    D(G(z)): 0.2423 / 0.0921
[1/5][1350/1583]        Loss_D: 0.8579  Loss_G: 4.6214  D(x): 0.9623    D(G(z)): 0.5089 / 0.0159
[1/5][1400/1583]        Loss_D: 0.4919  Loss_G: 2.7656  D(x): 0.8122    D(G(z)): 0.2147 / 0.0864
[1/5][1450/1583]        Loss_D: 0.4461  Loss_G: 3.0576  D(x): 0.8042    D(G(z)): 0.1798 / 0.0619
[1/5][1500/1583]        Loss_D: 0.7182  Loss_G: 3.7270  D(x): 0.8553    D(G(z)): 0.3713 / 0.0382
[1/5][1550/1583]        Loss_D: 0.6378  Loss_G: 3.7489  D(x): 0.8757    D(G(z)): 0.3523 / 0.0317
[2/5][0/1583]   Loss_D: 0.3965  Loss_G: 2.6262  D(x): 0.7941    D(G(z)): 0.1247 / 0.0963
[2/5][50/1583]  Loss_D: 0.6504  Loss_G: 3.9890  D(x): 0.9267    D(G(z)): 0.3865 / 0.0275
[2/5][100/1583] Loss_D: 0.6523  Loss_G: 3.8724  D(x): 0.8707    D(G(z)): 0.3613 / 0.0299
[2/5][150/1583] Loss_D: 0.7685  Loss_G: 3.9059  D(x): 0.9361    D(G(z)): 0.4534 / 0.0278
[2/5][200/1583] Loss_D: 0.6587  Loss_G: 1.9218  D(x): 0.6469    D(G(z)): 0.1291 / 0.1888
[2/5][250/1583] Loss_D: 0.6971  Loss_G: 2.2256  D(x): 0.6208    D(G(z)): 0.1226 / 0.1465
[2/5][300/1583] Loss_D: 0.5797  Loss_G: 2.4846  D(x): 0.7762    D(G(z)): 0.2434 / 0.1098
[2/5][350/1583] Loss_D: 0.4674  Loss_G: 1.8800  D(x): 0.8045    D(G(z)): 0.1903 / 0.1877
[2/5][400/1583] Loss_D: 0.6462  Loss_G: 1.9510  D(x): 0.7018    D(G(z)): 0.1935 / 0.1792
[2/5][450/1583] Loss_D: 0.9817  Loss_G: 4.2519  D(x): 0.9421    D(G(z)): 0.5381 / 0.0233
[2/5][500/1583] Loss_D: 0.7721  Loss_G: 1.0928  D(x): 0.5402    D(G(z)): 0.0316 / 0.3927
[2/5][550/1583] Loss_D: 0.6037  Loss_G: 2.6914  D(x): 0.7719    D(G(z)): 0.2504 / 0.0896
[2/5][600/1583] Loss_D: 1.4213  Loss_G: 5.4727  D(x): 0.9408    D(G(z)): 0.6792 / 0.0064
[2/5][650/1583] Loss_D: 0.7246  Loss_G: 1.7030  D(x): 0.6716    D(G(z)): 0.2184 / 0.2246
[2/5][700/1583] Loss_D: 0.6642  Loss_G: 3.3809  D(x): 0.8554    D(G(z)): 0.3438 / 0.0591
[2/5][750/1583] Loss_D: 0.6649  Loss_G: 2.0197  D(x): 0.7169    D(G(z)): 0.2333 / 0.1565
[2/5][800/1583] Loss_D: 0.4594  Loss_G: 2.6623  D(x): 0.8150    D(G(z)): 0.1930 / 0.0944
[2/5][850/1583] Loss_D: 1.1957  Loss_G: 3.1871  D(x): 0.7790    D(G(z)): 0.5576 / 0.0568
[2/5][900/1583] Loss_D: 0.6657  Loss_G: 1.5311  D(x): 0.7092    D(G(z)): 0.2122 / 0.2558
[2/5][950/1583] Loss_D: 0.6795  Loss_G: 1.4149  D(x): 0.6134    D(G(z)): 0.1195 / 0.2937
[2/5][1000/1583]        Loss_D: 0.5995  Loss_G: 2.1744  D(x): 0.7325    D(G(z)): 0.2054 / 0.1484
[2/5][1050/1583]        Loss_D: 0.6706  Loss_G: 1.6705  D(x): 0.6425    D(G(z)): 0.1414 / 0.2310
[2/5][1100/1583]        Loss_D: 1.2840  Loss_G: 4.4620  D(x): 0.9736    D(G(z)): 0.6601 / 0.0225
[2/5][1150/1583]        Loss_D: 0.7568  Loss_G: 3.1238  D(x): 0.8153    D(G(z)): 0.3717 / 0.0581
[2/5][1200/1583]        Loss_D: 0.6331  Loss_G: 1.9048  D(x): 0.6799    D(G(z)): 0.1604 / 0.1814
[2/5][1250/1583]        Loss_D: 0.5802  Loss_G: 2.4358  D(x): 0.7561    D(G(z)): 0.2194 / 0.1095
[2/5][1300/1583]        Loss_D: 0.9613  Loss_G: 2.3290  D(x): 0.7463    D(G(z)): 0.3952 / 0.1349
[2/5][1350/1583]        Loss_D: 0.5367  Loss_G: 1.7398  D(x): 0.7580    D(G(z)): 0.1898 / 0.2216
[2/5][1400/1583]        Loss_D: 0.7762  Loss_G: 3.6246  D(x): 0.9006    D(G(z)): 0.4378 / 0.0364
[2/5][1450/1583]        Loss_D: 0.7183  Loss_G: 4.0442  D(x): 0.8602    D(G(z)): 0.3857 / 0.0254
[2/5][1500/1583]        Loss_D: 0.5416  Loss_G: 2.0642  D(x): 0.7393    D(G(z)): 0.1758 / 0.1532
[2/5][1550/1583]        Loss_D: 0.5295  Loss_G: 1.7855  D(x): 0.6768    D(G(z)): 0.0886 / 0.2154
[3/5][0/1583]   Loss_D: 0.8635  Loss_G: 1.7508  D(x): 0.4918    D(G(z)): 0.0280 / 0.2154
[3/5][50/1583]  Loss_D: 0.8697  Loss_G: 0.7859  D(x): 0.5216    D(G(z)): 0.1124 / 0.4941
[3/5][100/1583] Loss_D: 0.8607  Loss_G: 4.5255  D(x): 0.9197    D(G(z)): 0.4973 / 0.0157
[3/5][150/1583] Loss_D: 0.4805  Loss_G: 2.3071  D(x): 0.7743    D(G(z)): 0.1742 / 0.1291
[3/5][200/1583] Loss_D: 0.4925  Loss_G: 2.6018  D(x): 0.7907    D(G(z)): 0.1970 / 0.0948
[3/5][250/1583] Loss_D: 0.7870  Loss_G: 3.3529  D(x): 0.8408    D(G(z)): 0.4050 / 0.0469
[3/5][300/1583] Loss_D: 0.5479  Loss_G: 1.7376  D(x): 0.7216    D(G(z)): 0.1592 / 0.2227
[3/5][350/1583] Loss_D: 0.8117  Loss_G: 3.4145  D(x): 0.9076    D(G(z)): 0.4685 / 0.0437
[3/5][400/1583] Loss_D: 0.4210  Loss_G: 2.3880  D(x): 0.7543    D(G(z)): 0.1047 / 0.1217
[3/5][450/1583] Loss_D: 1.5745  Loss_G: 0.2366  D(x): 0.2747    D(G(z)): 0.0361 / 0.8096
[3/5][500/1583] Loss_D: 0.7196  Loss_G: 2.1319  D(x): 0.7332    D(G(z)): 0.2935 / 0.1403
[3/5][550/1583] Loss_D: 0.5697  Loss_G: 2.6649  D(x): 0.8816    D(G(z)): 0.3210 / 0.0917
[3/5][600/1583] Loss_D: 0.7779  Loss_G: 1.2727  D(x): 0.5540    D(G(z)): 0.0855 / 0.3412
[3/5][650/1583] Loss_D: 0.4090  Loss_G: 2.6893  D(x): 0.8334    D(G(z)): 0.1835 / 0.0855
[3/5][700/1583] Loss_D: 0.8108  Loss_G: 3.8991  D(x): 0.9241    D(G(z)): 0.4716 / 0.0281
[3/5][750/1583] Loss_D: 0.9907  Loss_G: 4.7885  D(x): 0.9111    D(G(z)): 0.5402 / 0.0123
[3/5][800/1583] Loss_D: 0.4725  Loss_G: 2.3347  D(x): 0.7577    D(G(z)): 0.1400 / 0.1222
[3/5][850/1583] Loss_D: 1.5580  Loss_G: 4.9586  D(x): 0.8954    D(G(z)): 0.7085 / 0.0132
[3/5][900/1583] Loss_D: 0.5785  Loss_G: 1.6395  D(x): 0.6581    D(G(z)): 0.1003 / 0.2411
[3/5][950/1583] Loss_D: 0.6592  Loss_G: 1.0890  D(x): 0.5893    D(G(z)): 0.0451 / 0.3809
[3/5][1000/1583]        Loss_D: 0.7280  Loss_G: 3.5368  D(x): 0.8898    D(G(z)): 0.4176 / 0.0409
[3/5][1050/1583]        Loss_D: 0.7088  Loss_G: 3.4301  D(x): 0.8558    D(G(z)): 0.3845 / 0.0457
[3/5][1100/1583]        Loss_D: 0.5651  Loss_G: 2.1150  D(x): 0.7602    D(G(z)): 0.2127 / 0.1532
[3/5][1150/1583]        Loss_D: 0.5412  Loss_G: 1.7790  D(x): 0.6602    D(G(z)): 0.0801 / 0.2088
[3/5][1200/1583]        Loss_D: 1.2277  Loss_G: 1.1464  D(x): 0.4864    D(G(z)): 0.2915 / 0.3665
[3/5][1250/1583]        Loss_D: 0.7148  Loss_G: 1.3957  D(x): 0.5948    D(G(z)): 0.1076 / 0.2876
[3/5][1300/1583]        Loss_D: 1.0675  Loss_G: 1.3018  D(x): 0.4056    D(G(z)): 0.0310 / 0.3355
[3/5][1350/1583]        Loss_D: 0.8064  Loss_G: 0.7482  D(x): 0.5846    D(G(z)): 0.1453 / 0.5147
[3/5][1400/1583]        Loss_D: 0.6032  Loss_G: 3.0601  D(x): 0.8474    D(G(z)): 0.3189 / 0.0590
[3/5][1450/1583]        Loss_D: 0.5329  Loss_G: 2.8172  D(x): 0.8234    D(G(z)): 0.2567 / 0.0795
[3/5][1500/1583]        Loss_D: 0.9292  Loss_G: 3.5544  D(x): 0.8686    D(G(z)): 0.4887 / 0.0410
[3/5][1550/1583]        Loss_D: 0.5929  Loss_G: 2.9118  D(x): 0.8614    D(G(z)): 0.3239 / 0.0702
[4/5][0/1583]   Loss_D: 0.5564  Loss_G: 2.7516  D(x): 0.8716    D(G(z)): 0.3145 / 0.0799
[4/5][50/1583]  Loss_D: 1.0485  Loss_G: 0.6751  D(x): 0.4332    D(G(z)): 0.0675 / 0.5568
[4/5][100/1583] Loss_D: 0.6753  Loss_G: 1.4046  D(x): 0.6028    D(G(z)): 0.0882 / 0.2901
[4/5][150/1583] Loss_D: 0.5946  Loss_G: 1.7618  D(x): 0.6862    D(G(z)): 0.1488 / 0.2016
[4/5][200/1583] Loss_D: 0.4866  Loss_G: 2.2638  D(x): 0.7628    D(G(z)): 0.1633 / 0.1321
[4/5][250/1583] Loss_D: 0.7493  Loss_G: 1.0999  D(x): 0.5541    D(G(z)): 0.0659 / 0.3787
[4/5][300/1583] Loss_D: 1.0886  Loss_G: 4.6532  D(x): 0.9370    D(G(z)): 0.5811 / 0.0149
[4/5][350/1583] Loss_D: 0.6106  Loss_G: 1.9212  D(x): 0.6594    D(G(z)): 0.1322 / 0.1825
[4/5][400/1583] Loss_D: 0.5226  Loss_G: 2.9611  D(x): 0.8178    D(G(z)): 0.2378 / 0.0731
[4/5][450/1583] Loss_D: 1.0068  Loss_G: 1.3267  D(x): 0.4310    D(G(z)): 0.0375 / 0.3179
[4/5][500/1583] Loss_D: 3.1088  Loss_G: 0.1269  D(x): 0.0706    D(G(z)): 0.0061 / 0.8897
[4/5][550/1583] Loss_D: 1.7889  Loss_G: 0.4800  D(x): 0.2175    D(G(z)): 0.0143 / 0.6479
[4/5][600/1583] Loss_D: 0.6732  Loss_G: 3.5685  D(x): 0.8775    D(G(z)): 0.3879 / 0.0362
[4/5][650/1583] Loss_D: 0.5169  Loss_G: 2.1943  D(x): 0.7222    D(G(z)): 0.1349 / 0.1416
[4/5][700/1583] Loss_D: 0.4567  Loss_G: 2.4442  D(x): 0.7666    D(G(z)): 0.1410 / 0.1204
[4/5][750/1583] Loss_D: 0.5972  Loss_G: 2.2992  D(x): 0.6286    D(G(z)): 0.0670 / 0.1283
[4/5][800/1583] Loss_D: 0.5461  Loss_G: 1.9777  D(x): 0.7013    D(G(z)): 0.1318 / 0.1795
[4/5][850/1583] Loss_D: 0.6317  Loss_G: 2.2345  D(x): 0.6962    D(G(z)): 0.1854 / 0.1385
[4/5][900/1583] Loss_D: 0.6034  Loss_G: 3.2300  D(x): 0.8781    D(G(z)): 0.3448 / 0.0517
[4/5][950/1583] Loss_D: 0.6371  Loss_G: 2.7755  D(x): 0.8595    D(G(z)): 0.3357 / 0.0826
[4/5][1000/1583]        Loss_D: 0.6077  Loss_G: 3.3958  D(x): 0.9026    D(G(z)): 0.3604 / 0.0458
[4/5][1050/1583]        Loss_D: 0.5057  Loss_G: 3.2545  D(x): 0.8705    D(G(z)): 0.2691 / 0.0546
[4/5][1100/1583]        Loss_D: 0.4552  Loss_G: 2.0632  D(x): 0.7887    D(G(z)): 0.1704 / 0.1524
[4/5][1150/1583]        Loss_D: 0.9933  Loss_G: 1.0264  D(x): 0.4507    D(G(z)): 0.0636 / 0.4182
[4/5][1200/1583]        Loss_D: 0.5037  Loss_G: 1.9940  D(x): 0.6967    D(G(z)): 0.0959 / 0.1698
[4/5][1250/1583]        Loss_D: 0.4760  Loss_G: 2.5973  D(x): 0.8192    D(G(z)): 0.2164 / 0.0945
[4/5][1300/1583]        Loss_D: 1.0137  Loss_G: 3.8782  D(x): 0.9330    D(G(z)): 0.5405 / 0.0309
[4/5][1350/1583]        Loss_D: 0.9084  Loss_G: 3.1406  D(x): 0.7540    D(G(z)): 0.3980 / 0.0648
[4/5][1400/1583]        Loss_D: 0.6724  Loss_G: 4.1269  D(x): 0.9536    D(G(z)): 0.4234 / 0.0236
[4/5][1450/1583]        Loss_D: 0.6452  Loss_G: 3.5163  D(x): 0.8730    D(G(z)): 0.3555 / 0.0412
[4/5][1500/1583]        Loss_D: 0.8843  Loss_G: 1.4950  D(x): 0.5314    D(G(z)): 0.1035 / 0.2835
[4/5][1550/1583]        Loss_D: 2.3345  Loss_G: 1.0675  D(x): 0.1448    D(G(z)): 0.0228 / 0.4177 

plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(G_losses,label="G")
plt.plot(D_losses,label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.show() 

fig = plt.figure(figsize=(8,8))
plt.axis("off")
ims = [[plt.imshow(np.transpose(i,(1,2,0)), animated=True)] for i in img_list]
ani = animation.ArtistAnimation(fig, ims, interval=1000, repeat_delay=1000, blit=True)

HTML(ani.to_jshtml()) 

# License: BSD
# Author: Ghassen Hamrouni

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np

plt.ion()   # interactive mode 

<contextlib.ExitStack object at 0x7fc0914a7160> 

from six.moves import urllib
opener = urllib.request.build_opener()
opener.addheaders = [('User-agent', 'Mozilla/5.0')]
urllib.request.install_opener(opener)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Training dataset
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(root='.', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])), batch_size=64, shuffle=True, num_workers=4)
# Test dataset
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(root='.', train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])), batch_size=64, shuffle=True, num_workers=4) 

Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./MNIST/raw/train-images-idx3-ubyte.gz

  0%|          | 0/9912422 [00:00<?, ?it/s]
100%|##########| 9912422/9912422 [00:00<00:00, 367023704.91it/s]
Extracting ./MNIST/raw/train-images-idx3-ubyte.gz to ./MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./MNIST/raw/train-labels-idx1-ubyte.gz

  0%|          | 0/28881 [00:00<?, ?it/s]
100%|##########| 28881/28881 [00:00<00:00, 47653695.45it/s]
Extracting ./MNIST/raw/train-labels-idx1-ubyte.gz to ./MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./MNIST/raw/t10k-images-idx3-ubyte.gz

  0%|          | 0/1648877 [00:00<?, ?it/s]
100%|##########| 1648877/1648877 [00:00<00:00, 343101225.21it/s]
Extracting ./MNIST/raw/t10k-images-idx3-ubyte.gz to ./MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./MNIST/raw/t10k-labels-idx1-ubyte.gz

  0%|          | 0/4542 [00:00<?, ?it/s]
100%|##########| 4542/4542 [00:00<00:00, 48107395.88it/s]
Extracting ./MNIST/raw/t10k-labels-idx1-ubyte.gz to ./MNIST/raw 

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

        # Spatial transformer localization-network
        self.localization = nn.Sequential(
            nn.Conv2d(1, 8, kernel_size=7),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True),
            nn.Conv2d(8, 10, kernel_size=5),
            nn.MaxPool2d(2, stride=2),
            nn.ReLU(True)
        )

        # Regressor for the 3 * 2 affine matrix
        self.fc_loc = nn.Sequential(
            nn.Linear(10 * 3 * 3, 32),
            nn.ReLU(True),
            nn.Linear(32, 3 * 2)
        )

        # Initialize the weights/bias with identity transformation
        self.fc_loc[2].weight.data.zero_()
        self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))

    # Spatial transformer network forward function
    def stn(self, x):
        xs = self.localization(x)
        xs = xs.view(-1, 10 * 3 * 3)
        theta = self.fc_loc(xs)
        theta = theta.view(-1, 2, 3)

        grid = F.affine_grid(theta, x.size())
        x = F.grid_sample(x, grid)

        return x

    def forward(self, x):
        # transform the input
        x = self.stn(x)

        # Perform the usual forward pass
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = Net().to(device) 

optimizer = optim.SGD(model.parameters(), lr=0.01)

def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 500 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
#
# A simple test procedure to measure the STN performances on MNIST.
#

def test():
    with torch.no_grad():
        model.eval()
        test_loss = 0
        correct = 0
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)

            # sum up batch loss
            test_loss += F.nll_loss(output, target, size_average=False).item()
            # get the index of the max log-probability
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()

        test_loss /= len(test_loader.dataset)
        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'
              .format(test_loss, correct, len(test_loader.dataset),
                      100. * correct / len(test_loader.dataset))) 

def convert_image_np(inp):
  """Convert a Tensor to numpy image."""
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    return inp

# We want to visualize the output of the spatial transformers layer
# after the training, we visualize a batch of input images and
# the corresponding transformed batch using STN.

def visualize_stn():
    with torch.no_grad():
        # Get a batch of training data
        data = next(iter(test_loader))[0].to(device)

        input_tensor = data.cpu()
        transformed_input_tensor = model.stn(data).cpu()

        in_grid = convert_image_np(
            torchvision.utils.make_grid(input_tensor))

        out_grid = convert_image_np(
            torchvision.utils.make_grid(transformed_input_tensor))

        # Plot the results side-by-side
        f, axarr = plt.subplots(1, 2)
        axarr[0].imshow(in_grid)
        axarr[0].set_title('Dataset Images')

        axarr[1].imshow(out_grid)
        axarr[1].set_title('Transformed Images')

for epoch in range(1, 20 + 1):
    train(epoch)
    test()

# Visualize the STN transformation on some input batch
visualize_stn()

plt.ioff()
plt.show() 

/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/functional.py:4377: UserWarning:

Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0\. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.

/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/functional.py:4316: UserWarning:

Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0\. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.

Train Epoch: 1 [0/60000 (0%)]   Loss: 2.315648
Train Epoch: 1 [32000/60000 (53%)]      Loss: 1.051217
/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/_reduction.py:42: UserWarning:

size_average and reduce args will be deprecated, please use reduction='sum' instead.

Test set: Average loss: 0.2563, Accuracy: 9282/10000 (93%)

Train Epoch: 2 [0/60000 (0%)]   Loss: 0.544514
Train Epoch: 2 [32000/60000 (53%)]      Loss: 0.312879

Test set: Average loss: 0.1506, Accuracy: 9569/10000 (96%)

Train Epoch: 3 [0/60000 (0%)]   Loss: 0.408838
Train Epoch: 3 [32000/60000 (53%)]      Loss: 0.221301

Test set: Average loss: 0.1207, Accuracy: 9634/10000 (96%)

Train Epoch: 4 [0/60000 (0%)]   Loss: 0.400088
Train Epoch: 4 [32000/60000 (53%)]      Loss: 0.166533

Test set: Average loss: 0.1176, Accuracy: 9634/10000 (96%)

Train Epoch: 5 [0/60000 (0%)]   Loss: 0.274838
Train Epoch: 5 [32000/60000 (53%)]      Loss: 0.223936

Test set: Average loss: 0.2812, Accuracy: 9136/10000 (91%)

Train Epoch: 6 [0/60000 (0%)]   Loss: 0.411823
Train Epoch: 6 [32000/60000 (53%)]      Loss: 0.114000

Test set: Average loss: 0.0697, Accuracy: 9790/10000 (98%)

Train Epoch: 7 [0/60000 (0%)]   Loss: 0.066122
Train Epoch: 7 [32000/60000 (53%)]      Loss: 0.208773

Test set: Average loss: 0.0660, Accuracy: 9799/10000 (98%)

Train Epoch: 8 [0/60000 (0%)]   Loss: 0.201612
Train Epoch: 8 [32000/60000 (53%)]      Loss: 0.081877

Test set: Average loss: 0.0672, Accuracy: 9798/10000 (98%)

Train Epoch: 9 [0/60000 (0%)]   Loss: 0.077046
Train Epoch: 9 [32000/60000 (53%)]      Loss: 0.147858

Test set: Average loss: 0.0645, Accuracy: 9811/10000 (98%)

Train Epoch: 10 [0/60000 (0%)]  Loss: 0.086268
Train Epoch: 10 [32000/60000 (53%)]     Loss: 0.185868

Test set: Average loss: 0.0678, Accuracy: 9794/10000 (98%)

Train Epoch: 11 [0/60000 (0%)]  Loss: 0.138696
Train Epoch: 11 [32000/60000 (53%)]     Loss: 0.119381

Test set: Average loss: 0.0663, Accuracy: 9795/10000 (98%)

Train Epoch: 12 [0/60000 (0%)]  Loss: 0.145220
Train Epoch: 12 [32000/60000 (53%)]     Loss: 0.204023

Test set: Average loss: 0.0592, Accuracy: 9808/10000 (98%)

Train Epoch: 13 [0/60000 (0%)]  Loss: 0.118743
Train Epoch: 13 [32000/60000 (53%)]     Loss: 0.100721

Test set: Average loss: 0.0643, Accuracy: 9801/10000 (98%)

Train Epoch: 14 [0/60000 (0%)]  Loss: 0.066341
Train Epoch: 14 [32000/60000 (53%)]     Loss: 0.107528

Test set: Average loss: 0.0551, Accuracy: 9838/10000 (98%)

Train Epoch: 15 [0/60000 (0%)]  Loss: 0.022679
Train Epoch: 15 [32000/60000 (53%)]     Loss: 0.055676

Test set: Average loss: 0.0474, Accuracy: 9862/10000 (99%)

Train Epoch: 16 [0/60000 (0%)]  Loss: 0.102644
Train Epoch: 16 [32000/60000 (53%)]     Loss: 0.165537

Test set: Average loss: 0.0574, Accuracy: 9839/10000 (98%)

Train Epoch: 17 [0/60000 (0%)]  Loss: 0.280918
Train Epoch: 17 [32000/60000 (53%)]     Loss: 0.206559

Test set: Average loss: 0.0533, Accuracy: 9846/10000 (98%)

Train Epoch: 18 [0/60000 (0%)]  Loss: 0.052316
Train Epoch: 18 [32000/60000 (53%)]     Loss: 0.082710

Test set: Average loss: 0.0484, Accuracy: 9865/10000 (99%)

Train Epoch: 19 [0/60000 (0%)]  Loss: 0.083889
Train Epoch: 19 [32000/60000 (53%)]     Loss: 0.121432

Test set: Average loss: 0.0522, Accuracy: 9839/10000 (98%)

Train Epoch: 20 [0/60000 (0%)]  Loss: 0.067540
Train Epoch: 20 [32000/60000 (53%)]     Loss: 0.024880

Test set: Average loss: 0.0868, Accuracy: 9773/10000 (98%) 

pip install torch torchvision timm pandas requests 

!pip install timm pandas requests 

from PIL import Image
import torch
import timm
import requests
import torchvision.transforms as transforms
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD

print(torch.__version__)
# should be 1.8.0

model = torch.hub.load('facebookresearch/deit:main', 'deit_base_patch16_224', pretrained=True)
model.eval()

transform = transforms.Compose([
    transforms.Resize(256, interpolation=3),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD),
])

img = Image.open(requests.get("https://raw.githubusercontent.com/pytorch/ios-demo-app/master/HelloWorld/HelloWorld/HelloWorld/image.png", stream=True).raw)
img = transform(img)[None,]
out = model(img)
clsidx = torch.argmax(out)
print(clsidx.item()) 

2.2.0+cu121
Downloading: "https://github.com/facebookresearch/deit/zipball/main" to /var/lib/jenkins/.cache/torch/hub/main.zip
/var/lib/jenkins/.cache/torch/hub/facebookresearch_deit_main/models.py:63: UserWarning:

Overwriting deit_tiny_patch16_224 in registry with models.deit_tiny_patch16_224\. This is because the name being registered conflicts with an existing name. Please check if this is not expected.

/var/lib/jenkins/.cache/torch/hub/facebookresearch_deit_main/models.py:78: UserWarning:

Overwriting deit_small_patch16_224 in registry with models.deit_small_patch16_224\. This is because the name being registered conflicts with an existing name. Please check if this is not expected.

/var/lib/jenkins/.cache/torch/hub/facebookresearch_deit_main/models.py:93: UserWarning:

Overwriting deit_base_patch16_224 in registry with models.deit_base_patch16_224\. This is because the name being registered conflicts with an existing name. Please check if this is not expected.

/var/lib/jenkins/.cache/torch/hub/facebookresearch_deit_main/models.py:108: UserWarning:

Overwriting deit_tiny_distilled_patch16_224 in registry with models.deit_tiny_distilled_patch16_224\. This is because the name being registered conflicts with an existing name. Please check if this is not expected.

/var/lib/jenkins/.cache/torch/hub/facebookresearch_deit_main/models.py:123: UserWarning:

Overwriting deit_small_distilled_patch16_224 in registry with models.deit_small_distilled_patch16_224\. This is because the name being registered conflicts with an existing name. Please check if this is not expected.

/var/lib/jenkins/.cache/torch/hub/facebookresearch_deit_main/models.py:138: UserWarning:

Overwriting deit_base_distilled_patch16_224 in registry with models.deit_base_distilled_patch16_224\. This is because the name being registered conflicts with an existing name. Please check if this is not expected.

/var/lib/jenkins/.cache/torch/hub/facebookresearch_deit_main/models.py:153: UserWarning:

Overwriting deit_base_patch16_384 in registry with models.deit_base_patch16_384\. This is because the name being registered conflicts with an existing name. Please check if this is not expected.

/var/lib/jenkins/.cache/torch/hub/facebookresearch_deit_main/models.py:168: UserWarning:

Overwriting deit_base_distilled_patch16_384 in registry with models.deit_base_distilled_patch16_384\. This is because the name being registered conflicts with an existing name. Please check if this is not expected.

Downloading: "https://dl.fbaipublicfiles.com/deit/deit_base_patch16_224-b5f2ef4d.pth" to /var/lib/jenkins/.cache/torch/hub/checkpoints/deit_base_patch16_224-b5f2ef4d.pth

  0%|          | 0.00/330M [00:00<?, ?B/s]
  4%|3         | 12.4M/330M [00:00<00:02, 130MB/s]
  7%|7         | 24.7M/330M [00:00<00:02, 110MB/s]
 11%|#1        | 36.8M/330M [00:00<00:02, 117MB/s]
 15%|#4        | 49.2M/330M [00:00<00:02, 121MB/s]
 19%|#8        | 62.2M/330M [00:00<00:02, 127MB/s]
 23%|##3       | 76.7M/330M [00:00<00:01, 135MB/s]
 27%|##7       | 90.6M/330M [00:00<00:01, 139MB/s]
 32%|###1      | 106M/330M [00:00<00:01, 144MB/s]
 36%|###6      | 119M/330M [00:00<00:01, 125MB/s]
 40%|###9      | 132M/330M [00:01<00:01, 122MB/s]
 45%|####4     | 147M/330M [00:01<00:01, 132MB/s]
 49%|####8     | 162M/330M [00:01<00:01, 138MB/s]
 53%|#####3    | 176M/330M [00:01<00:01, 142MB/s]
 58%|#####7    | 190M/330M [00:01<00:01, 144MB/s]
 62%|######2   | 205M/330M [00:01<00:00, 147MB/s]
 67%|######6   | 220M/330M [00:01<00:00, 149MB/s]
 71%|#######   | 234M/330M [00:01<00:00, 148MB/s]
 76%|#######5  | 250M/330M [00:01<00:00, 155MB/s]
 81%|########1 | 268M/330M [00:01<00:00, 162MB/s]
 86%|########6 | 285M/330M [00:02<00:00, 168MB/s]
 91%|#########1| 302M/330M [00:02<00:00, 172MB/s]
 97%|#########6| 319M/330M [00:02<00:00, 175MB/s]
100%|##########| 330M/330M [00:02<00:00, 147MB/s]
269 

model = torch.hub.load('facebookresearch/deit:main', 'deit_base_patch16_224', pretrained=True)
model.eval()
scripted_model = torch.jit.script(model)
scripted_model.save("fbdeit_scripted.pt") 

Using cache found in /var/lib/jenkins/.cache/torch/hub/facebookresearch_deit_main 

# Use 'x86' for server inference (the old 'fbgemm' is still available but 'x86' is the recommended default) and ``qnnpack`` for mobile inference.
backend = "x86" # replaced with ``qnnpack`` causing much worse inference speed for quantized model on this notebook
model.qconfig = torch.quantization.get_default_qconfig(backend)
torch.backends.quantized.engine = backend

quantized_model = torch.quantization.quantize_dynamic(model, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)
scripted_quantized_model = torch.jit.script(quantized_model)
scripted_quantized_model.save("fbdeit_scripted_quantized.pt") 

/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/ao/quantization/observer.py:220: UserWarning:

Please use quant_min and quant_max to specify the range for observers.                     reduce_range will be deprecated in a future release of PyTorch. 

out = scripted_quantized_model(img)
clsidx = torch.argmax(out)
print(clsidx.item())
# The same output 269 should be printed 

269 

from torch.utils.mobile_optimizer import optimize_for_mobile
optimized_scripted_quantized_model = optimize_for_mobile(scripted_quantized_model)
optimized_scripted_quantized_model.save("fbdeit_optimized_scripted_quantized.pt") 

out = optimized_scripted_quantized_model(img)
clsidx = torch.argmax(out)
print(clsidx.item())
# Again, the same output 269 should be printed 

269 

optimized_scripted_quantized_model._save_for_lite_interpreter("fbdeit_optimized_scripted_quantized_lite.ptl")
ptl = torch.jit.load("fbdeit_optimized_scripted_quantized_lite.ptl") 

with torch.autograd.profiler.profile(use_cuda=False) as prof1:
    out = model(img)
with torch.autograd.profiler.profile(use_cuda=False) as prof2:
    out = scripted_model(img)
with torch.autograd.profiler.profile(use_cuda=False) as prof3:
    out = scripted_quantized_model(img)
with torch.autograd.profiler.profile(use_cuda=False) as prof4:
    out = optimized_scripted_quantized_model(img)
with torch.autograd.profiler.profile(use_cuda=False) as prof5:
    out = ptl(img)

print("original model: {:.2f}ms".format(prof1.self_cpu_time_total/1000))
print("scripted model: {:.2f}ms".format(prof2.self_cpu_time_total/1000))
print("scripted & quantized model: {:.2f}ms".format(prof3.self_cpu_time_total/1000))
print("scripted & quantized & optimized model: {:.2f}ms".format(prof4.self_cpu_time_total/1000))
print("lite model: {:.2f}ms".format(prof5.self_cpu_time_total/1000)) 

original model: 123.27ms
scripted model: 111.89ms
scripted & quantized model: 129.99ms
scripted & quantized & optimized model: 129.94ms
lite model: 120.00ms 

original  model:  1236.69ms
scripted  model:  1226.72ms
scripted  &  quantized  model:  593.19ms
scripted  &  quantized  &  optimized  model:  598.01ms
lite  model:  600.72ms 

import pandas as pd
import numpy as np

df = pd.DataFrame({'Model': ['original model','scripted model', 'scripted & quantized model', 'scripted & quantized & optimized model', 'lite model']})
df = pd.concat([df, pd.DataFrame([
    ["{:.2f}ms".format(prof1.self_cpu_time_total/1000), "0%"],
    ["{:.2f}ms".format(prof2.self_cpu_time_total/1000),
     "{:.2f}%".format((prof1.self_cpu_time_total-prof2.self_cpu_time_total)/prof1.self_cpu_time_total*100)],
    ["{:.2f}ms".format(prof3.self_cpu_time_total/1000),
     "{:.2f}%".format((prof1.self_cpu_time_total-prof3.self_cpu_time_total)/prof1.self_cpu_time_total*100)],
    ["{:.2f}ms".format(prof4.self_cpu_time_total/1000),
     "{:.2f}%".format((prof1.self_cpu_time_total-prof4.self_cpu_time_total)/prof1.self_cpu_time_total*100)],
    ["{:.2f}ms".format(prof5.self_cpu_time_total/1000),
     "{:.2f}%".format((prof1.self_cpu_time_total-prof5.self_cpu_time_total)/prof1.self_cpu_time_total*100)]],
    columns=['Inference Time', 'Reduction'])], axis=1)

print(df)

"""
 Model                             Inference Time    Reduction
0   original model                             1236.69ms           0%
1   scripted model                             1226.72ms        0.81%
2   scripted & quantized model                  593.19ms       52.03%
3   scripted & quantized & optimized model      598.01ms       51.64%
4   lite model                                  600.72ms       51.43%
""" 

 Model  ... Reduction
0                          original model  ...        0%
1                          scripted model  ...     9.23%
2              scripted & quantized model  ...    -5.45%
3  scripted & quantized & optimized model  ...    -5.41%
4                              lite model  ...     2.65%

[5 rows x 3 columns]

'\n        Model                             Inference Time    Reduction\n0\toriginal model                             1236.69ms           0%\n1\tscripted model                             1226.72ms        0.81%\n2\tscripted & quantized model                  593.19ms       52.03%\n3\tscripted & quantized & optimized model      598.01ms       51.64%\n4\tlite model                                  600.72ms       51.43%\n' 

"""Import modules required to run the Jupyter notebook."""
from __future__ import annotations

# Configure logging
import logging
import warnings
if logging.getLogger().hasHandlers():
    logging.getLogger().handlers.clear()
warnings.filterwarnings("ignore", message=".*The 'nopython' keyword.*")

# Downloading data and files
import shutil
from pathlib import Path
from zipfile import ZipFile

# Data processing and visualization
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib import cm
import PIL
import contextlib
import io
from sklearn.metrics import accuracy_score, confusion_matrix

# TIAToolbox for WSI loading and processing
from tiatoolbox import logger
from tiatoolbox.models.architecture import vanilla
from tiatoolbox.models.engine.patch_predictor import (
    IOPatchPredictorConfig,
    PatchPredictor,
)
from tiatoolbox.utils.misc import download_data, grab_files_from_dir
from tiatoolbox.utils.visualization import overlay_prediction_mask
from tiatoolbox.wsicore.wsireader import WSIReader

# Torch-related
import torch
from torchvision import transforms

# Configure plotting
mpl.rcParams["figure.dpi"] = 160  # for high resolution figure in notebook
mpl.rcParams["figure.facecolor"] = "white"  # To make sure text is visible in dark mode

# If you are not using GPU, change ON_GPU to False
ON_GPU = True

# Function to suppress console output for overly verbose code blocks
def suppress_console_output():
    return contextlib.redirect_stderr(io.StringIO()) 

warnings.filterwarnings("ignore")
global_save_dir = Path("./tmp/")

def rmdir(dir_path: str | Path) -> None:
  """Helper function to delete directory."""
    if Path(dir_path).is_dir():
        shutil.rmtree(dir_path)
        logger.info("Removing directory %s", dir_path)

rmdir(global_save_dir)  # remove  directory if it exists from previous runs
global_save_dir.mkdir()
logger.info("Creating new directory %s", global_save_dir) 

wsi_path = global_save_dir / "sample_wsi.svs"
patches_path = global_save_dir / "kather100k-validation-sample.zip"
weights_path = global_save_dir / "resnet18-kather100k.pth"

logger.info("Download has started. Please wait...")

# Downloading and unzip a sample whole-slide image
download_data(
    "https://tiatoolbox.dcs.warwick.ac.uk/sample_wsis/TCGA-3L-AA1B-01Z-00-DX1.8923A151-A690-40B7-9E5A-FCBEDFC2394F.svs",
    wsi_path,
)

# Download and unzip a sample of the validation set used to train the Kather 100K dataset
download_data(
    "https://tiatoolbox.dcs.warwick.ac.uk/datasets/kather100k-validation-sample.zip",
    patches_path,
)
with ZipFile(patches_path, "r") as zipfile:
    zipfile.extractall(path=global_save_dir)

# Download pretrained model weights for WSI classification using ResNet18 architecture
download_data(
    "https://tiatoolbox.dcs.warwick.ac.uk/models/pc/resnet18-kather100k.pth",
    weights_path,
)

logger.info("Download is complete.") 

# Read the patch data and create a list of patches and a list of corresponding labels
dataset_path = global_save_dir / "kather100k-validation-sample"

# Set the path to the dataset
image_ext = ".tif"  # file extension of each image

# Obtain the mapping between the label ID and the class name
label_dict = {
    "BACK": 0, # Background (empty glass region)
    "NORM": 1, # Normal colon mucosa
    "DEB": 2,  # Debris
    "TUM": 3,  # Colorectal adenocarcinoma epithelium
    "ADI": 4,  # Adipose
    "MUC": 5,  # Mucus
    "MUS": 6,  # Smooth muscle
    "STR": 7,  # Cancer-associated stroma
    "LYM": 8,  # Lymphocytes
}

class_names = list(label_dict.keys())
class_labels = list(label_dict.values())

# Generate a list of patches and generate the label from the filename
patch_list = []
label_list = []
for class_name, label in label_dict.items():
    dataset_class_path = dataset_path / class_name
    patch_list_single_class = grab_files_from_dir(
        dataset_class_path,
        file_types="*" + image_ext,
    )
    patch_list.extend(patch_list_single_class)
    label_list.extend([label] * len(patch_list_single_class))

# Show some dataset statistics
plt.bar(class_names, [label_list.count(label) for label in class_labels])
plt.xlabel("Patch types")
plt.ylabel("Number of patches")

# Count the number of examples per class
for class_name, label in label_dict.items():
    logger.info(
        "Class ID: %d -- Class Name: %s -- Number of images: %d",
        label,
        class_name,
        label_list.count(label),
    )

# Overall dataset statistics
logger.info("Total number of patches: %d", (len(patch_list))) 

|2023-11-14|13:15:59.299| [INFO] Class ID: 0 -- Class Name: BACK -- Number of images: 211
|2023-11-14|13:15:59.299| [INFO] Class ID: 1 -- Class Name: NORM -- Number of images: 176
|2023-11-14|13:15:59.299| [INFO] Class ID: 2 -- Class Name: DEB -- Number of images: 230
|2023-11-14|13:15:59.299| [INFO] Class ID: 3 -- Class Name: TUM -- Number of images: 286
|2023-11-14|13:15:59.299| [INFO] Class ID: 4 -- Class Name: ADI -- Number of images: 208
|2023-11-14|13:15:59.299| [INFO] Class ID: 5 -- Class Name: MUC -- Number of images: 178
|2023-11-14|13:15:59.299| [INFO] Class ID: 6 -- Class Name: MUS -- Number of images: 270
|2023-11-14|13:15:59.299| [INFO] Class ID: 7 -- Class Name: STR -- Number of images: 209
|2023-11-14|13:15:59.299| [INFO] Class ID: 8 -- Class Name: LYM -- Number of images: 232
|2023-11-14|13:15:59.299| [INFO] Total number of patches: 2000 

# Importing a pretrained PyTorch model from TIAToolbox
predictor = PatchPredictor(pretrained_model='resnet18-kather100k', batch_size=32)

# Users can load any PyTorch model architecture instead using the following script
model = vanilla.CNNModel(backbone="resnet18", num_classes=9) # Importing model from torchvision.models.resnet18
model.load_state_dict(torch.load(weights_path, map_location="cpu"), strict=True)
def preproc_func(img):
    img = PIL.Image.fromarray(img)
    img = transforms.ToTensor()(img)
    return img.permute(1, 2, 0)
model.preproc_func = preproc_func
predictor = PatchPredictor(model=model, batch_size=32) 

with suppress_console_output():
    output = predictor.predict(imgs=patch_list, mode="patch", on_gpu=ON_GPU)

acc = accuracy_score(label_list, output["predictions"])
logger.info("Classification accuracy: %f", acc)

# Creating and visualizing the confusion matrix for patch classification results
conf = confusion_matrix(label_list, output["predictions"], normalize="true")
df_cm = pd.DataFrame(conf, index=class_names, columns=class_names)
df_cm 

|2023-11-14|13:16:03.215| [INFO] Classification accuracy: 0.993000 

wsi_ioconfig = IOPatchPredictorConfig(
    input_resolutions=[{"units": "mpp", "resolution": 0.5}],
    patch_input_shape=[224, 224],
    stride_shape=[224, 224],
) 

with suppress_console_output():
    wsi_output = predictor.predict(
        imgs=[wsi_path],
        masks=None,
        mode="wsi",
        merge_predictions=False,
        ioconfig=wsi_ioconfig,
        return_probabilities=True,
        save_dir=global_save_dir / "wsi_predictions",
        on_gpu=ON_GPU,
    ) 

overview_resolution = (
    4  # the resolution in which we desire to merge and visualize the patch predictions
)
# the unit of the `resolution` parameter. Can be "power", "level", "mpp", or "baseline"
overview_unit = "mpp"
wsi = WSIReader.open(wsi_path)
wsi_overview = wsi.slide_thumbnail(resolution=overview_resolution, units=overview_unit)
plt.figure(), plt.imshow(wsi_overview)
plt.axis("off") 

# Visualization of whole-slide image patch-level prediction
# first set up a label to color mapping
label_color_dict = {}
label_color_dict[0] = ("empty", (0, 0, 0))
colors = cm.get_cmap("Set1").colors
for class_name, label in label_dict.items():
    label_color_dict[label + 1] = (class_name, 255 * np.array(colors[label]))

pred_map = predictor.merge_predictions(
    wsi_path,
    wsi_output[0],
    resolution=overview_resolution,
    units=overview_unit,
)
overlay = overlay_prediction_mask(
    wsi_overview,
    pred_map,
    alpha=0.5,
    label_info=label_color_dict,
    return_ax=True,
)
plt.show() 

# Import some extra modules
import histoencoder.functional as F
import torch.nn as nn

from tiatoolbox.models.engine.semantic_segmentor import DeepFeatureExtractor, IOSegmentorConfig
from tiatoolbox.models.models_abc import ModelABC
import umap 

class HistoEncWrapper(ModelABC):
  """Wrapper for HistoEnc model that conforms to tiatoolbox ModelABC interface."""

    def __init__(self: HistoEncWrapper, encoder) -> None:
        super().__init__()
        self.feat_extract = encoder

    def forward(self: HistoEncWrapper, imgs: torch.Tensor) -> torch.Tensor:
  """Pass input data through the model.

 Args:
 imgs (torch.Tensor):
 Model input.

 """
        out = F.extract_features(self.feat_extract, imgs, num_blocks=2, avg_pool=True)
        return out

    @staticmethod
    def infer_batch(
        model: nn.Module,
        batch_data: torch.Tensor,
        *,
        on_gpu: bool,
    ) -> list[np.ndarray]:
  """Run inference on an input batch.

 Contains logic for forward operation as well as i/o aggregation.

 Args:
 model (nn.Module):
 PyTorch defined model.
 batch_data (torch.Tensor):
 A batch of data generated by
 `torch.utils.data.DataLoader`.
 on_gpu (bool):
 Whether to run inference on a GPU.

 """
        img_patches_device = batch_data.to('cuda') if on_gpu else batch_data
        model.eval()
        # Do not compute the gradient (not training)
        with torch.inference_mode():
            output = model(img_patches_device)
        return [output.cpu().numpy()] 

# create the model
encoder = F.create_encoder("prostate_medium")
model = HistoEncWrapper(encoder)

# set the pre-processing function
norm=transforms.Normalize(mean=[0.662, 0.446, 0.605],std=[0.169, 0.190, 0.155])
trans = [
    transforms.ToTensor(),
    norm,
]
model.preproc_func = transforms.Compose(trans)

wsi_ioconfig = IOSegmentorConfig(
    input_resolutions=[{"units": "mpp", "resolution": 0.5}],
    patch_input_shape=[224, 224],
    output_resolutions=[{"units": "mpp", "resolution": 0.5}],
    patch_output_shape=[224, 224],
    stride_shape=[224, 224],
) 

# create the feature extractor and run it on the WSI
extractor = DeepFeatureExtractor(model=model, auto_generate_mask=True, batch_size=32, num_loader_workers=4, num_postproc_workers=4)
with suppress_console_output():
    out = extractor.predict(imgs=[wsi_path], mode="wsi", ioconfig=wsi_ioconfig, save_dir=global_save_dir / "wsi_features",) 

# First we define a function to calculate the umap reduction
def umap_reducer(x, dims=3, nns=10):
  """UMAP reduction of the input data."""
    reducer = umap.UMAP(n_neighbors=nns, n_components=dims, metric="manhattan", spread=0.5, random_state=2)
    reduced = reducer.fit_transform(x)
    reduced -= reduced.min(axis=0)
    reduced /= reduced.max(axis=0)
    return reduced

# load the features output by our feature extractor
pos = np.load(global_save_dir / "wsi_features" / "0.position.npy")
feats = np.load(global_save_dir / "wsi_features" / "0.features.0.npy")
pos = pos / 8 # as we extracted at 0.5mpp, and we are overlaying on a thumbnail at 4mpp

# reduce the features into 3 dimensional (rgb) space
reduced = umap_reducer(feats)

# plot the prediction map the classifier again
overlay = overlay_prediction_mask(
    wsi_overview,
    pred_map,
    alpha=0.5,
    label_info=label_color_dict,
    return_ax=True,
)

# plot the feature map reduction
plt.figure()
plt.imshow(wsi_overview)
plt.scatter(pos[:,0], pos[:,1], c=reduced, s=1, alpha=0.5)
plt.axis("off")
plt.title("UMAP reduction of HistoEnc features")
plt.show()