<jupyter_start><jupyter_text>This notebook will train a cat-dog classification model, using the pretrianed resnet18 as the feature extractor.The trainig code has been borrowd from [Pytorch vision transfer learning tutorial](https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html), and the dataset is from taken from Kaggle. You can read more on how to use Kaggle datasets on Google Colab  [here](https://www.kaggle.com/general/74235).<jupyter_code>! pip install -q kaggle
# Upload kaggle.json file with kaggle API token here
from google.colab import files
uploaded = files.upload()
! mkdir ~/.kaggle
! cp kaggle.json ~/.kaggle/
! chmod 600 ~/.kaggle/kaggle.json
!kaggle competitions download -c dogs-vs-cats
!unzip train.zip
!unzip test1.zip
import os
train_files = os.listdir('./train')

!mkdir dog-cat-dataset
!mkdir dog-cat-dataset/dogs
!mkdir dog-cat-dataset/cats

!pwd<jupyter_output><empty_output><jupyter_text>Splitting the cat and dog images in the dataset.<jupyter_code>import shutil
base_dir='/content/train'
for filename in train_files:

  category = filename.split('.')[0]
  if category=='dog':
    shutil.copyfile(os.path.join(base_dir,filename), os.path.join('/content/dog-cat-dataset/dogs',filename))
  elif category=='cat':
    shutil.copyfile(os.path.join(base_dir,filename), os.path.join('/content/dog-cat-dataset/cats',filename))

%cd /content/dog-cat-dataset/cats
!ls | wc -l 
%cd -
!mkdir ./dog-cat-dataset/test
!mkdir ./dog-cat-dataset/test/dog
!mkdir ./dog-cat-dataset/test/cat
!mkdir ./dog-cat-dataset/val
!mkdir ./dog-cat-dataset/val/dog
!mkdir ./dog-cat-dataset/val/cat
!mkdir ./dog-cat-dataset/train
!mkdir ./dog-cat-dataset/train/dog
!mkdir ./dog-cat-dataset/train/cat
%cd ./dog-cat-dataset/dogs/
!ls | wc -l<jupyter_output><empty_output><jupyter_text>Here, we split the dog and cat images in two train/val/test sets.<jupyter_code>#% cd dog-cat-dataset/dogs
% ls | shuf -n 2250 | xargs -i mv {} ../val/dog
% ls | shuf -n 2250 | xargs -i mv {} ../test/dog
% ls | shuf -n 8000 | xargs -i mv {} ../train/dog
% cd ../cats
% ls | shuf -n 2250 | xargs -i mv {} ../val/cat
% ls | shuf -n 2250 | xargs -i mv {} ../test/cat
% ls | shuf -n 8000 | xargs -i mv {} ../train/cat

%cd ../../
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data
import torch.nn.functional as F
import torchvision
import torchvision.models as models
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
from torch.optim import lr_scheduler
from torchvision import datasets, models, transforms
import os<jupyter_output><empty_output><jupyter_text>Setting the dataset, data loader and the image transforms for different sets.<jupyter_code>data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'test': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

data_dir = '/content/dog-cat-dataset/'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val', 'test']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=32,
                                             shuffle=True, num_workers=2)
              for x in ['train', 'val','test']}<jupyter_output><empty_output><jupyter_text>Loading the pretrained resnet18 from torch-hub and changing the last layer to the number of features to 512, adding a dropout layer and output two classes of cat and dog.<jupyter_code>pretrained_model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)

for name, param in pretrained_model.named_parameters():
    if("bn" not in name):
        param.requires_grad = False
        
num_ftrs = pretrained_model.fc.in_features

pretrained_model.fc = nn.Sequential(nn.Linear(pretrained_model.fc.in_features,512),
                                  nn.ReLU(),
                                  nn.Dropout(),
                                  nn.Linear(512, 2))<jupyter_output><empty_output><jupyter_text>Setting the optimizer and learning rate scheduler.<jupyter_code>optimizer = optim.Adam(pretrained_model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)<jupyter_output><empty_output><jupyter_text>Define the training loop.<jupyter_code>def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model
import time
import copy
if torch.cuda.is_available():
    device = torch.device("cuda") 
else:
    device = torch.device("cpu")
pretrained_model.to(device)
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes
print(class_names)
print(dataset_sizes)<jupyter_output><empty_output><jupyter_text>Run the training for 2 epochs.<jupyter_code>pretrained_model = train_model(pretrained_model, criterion, optimizer, exp_lr_scheduler,
                       num_epochs=2)<jupyter_output><empty_output><jupyter_text>Testing the model pefromance on the test set.<jupyter_code>def test_model(model):
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in dataloaders['test']:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print('correct: {:d}  total: {:d}'.format(correct, total))
    print('accuracy = {:f}'.format(correct / total))
test_model(pretrained_model)
torch.save(pretrained_model.state_dict(), "./cat_dog_classification.pth")<jupyter_output><empty_output><jupyter_text>Load the trained/saved model for sanity check and testing on the test dataset.<jupyter_code>pretrained_model = torch.hub.load('pytorch/vision', 'resnet18')
pretrained_model.fc = nn.Sequential(nn.Linear(pretrained_model.fc.in_features,512),nn.ReLU(), nn.Dropout(), nn.Linear(512, 2))
pretrained_model.load_state_dict(torch.load('./cat_dog_classification.pth'))
pretrained_model.eval()
pretrained_model.to(device)
test_model(pretrained_model)<jupyter_output><empty_output><jupyter_text>Running some prediction tests.<jupyter_code>def prediction(model, filename):
    labels = class_names
    img = Image.open(filename)
    img = data_transforms['test'](img)
    img = img.unsqueeze(0)
    prediction = model(img.to(device))
    print(prediction)
    prediction = prediction.argmax()
    print(labels)
    print(labels[prediction])
prediction(pretrained_model, '/content/dog-cat-dataset/test/cat/cat.100.jpg')
prediction(pretrained_model, '/content/dog-cat-dataset/test/dog/dog.10033.jpg')<jupyter_output><empty_output><jupyter_text>Downloading the trained model.<jupyter_code>from google.colab import files
files.download('./cat_dog_classification.pth')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12447279086513734, 'fraction_numerical': 0.02005966464355519, 'mean_word_length': 3.04914618908788, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 50, 'https://': 2, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '304626', 'license': 'Apache-2.0', 'n_tokens_mistral': 3308, 'n_tokens_neox': 3069, 'n_words': 700, 'path': 'examples/Workflows/dog_breed_classification/cat_dog_classification.ipynb', 'repo_name': 'brianjo/serve', 'size': '9721'}