n0obcoder

import torch
import torch.nn as nn
from torchvision import models

# Defining the model
model = models.resnet34(pretrained = True)

# The original architecture of resnet34 has 1000 neurons(corresponding to 1000 classes on which it was originally trained on) in the final layer.
# So we need to change the final layer according to the number of classes that we have in our dataset
print('model.fc before: ', model.fc)

n0obcoder

# Now let's have a look at the requires_grad attributes for all the parameters

for name, param in model.named_parameters():
    print('name: {} has requires_grad: {}'.format(name, param.requires_grad))

n0obcoder

for name, module in model.named_children():
    print('name: ', name)

n0obcoder

#  We would freeze all but the last few layers (layer4 and fc)

for name, param in model.named_parameters():
    if ('layer4' in name) or ('fc' in name):
        param.requires_grad = True
    else:
        param.requires_grad = False

n0obcoder

for name, param in model.named_parameters():
    print('name: {} has requires_grad: {}'.format(name, param.requires_grad))

n0obcoder

import torch.optim as optim
from torch.optim import lr_scheduler

# Now we define the Loss Function
loss_fn = nn.CrossEntropyLoss()

# Define the optimizer
lr = 0.00001
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = lr)

n0obcoder

# Defining a trainer function which would train and validate the model

def trainer(dataloader_dict, model, loss_fn, optimizer, epochs = 1, log_interval = 1):
    print('Training started...')
    train_losses = []
    val_losses = []
    batch_train_losses = []
    batch_val_losses = []
    for epoch in range(epochs):
        print('epoch >>> {}/{}'.format(epoch + 1, epochs))

n0obcoder

loader = {'train': train_loader, 'val': val_loader}

epochs = 5
log_interval = 2

# Let's train the model for 5 epochs !
train_losses, val_losses, batch_train_losses, batch_val_losses = trainer(loader, model, loss_fn, optimizer, epochs = epochs, log_interval = log_interval)

# Ploting the epoch losses
plt.plot(train_losses)

n0obcoder

# We will now freeze the 'layer4' and train just the 'fc' layer of the model for 2 more epochs

for name, param in model.named_parameters():
    if 'layer4' in name:
        param.requires_grad = False # layer4 parameters would not get trained now
        

# Define the new learning rate and the new optimizer which would contain only the parameters with requires_grad = True 
lr = 0.0003
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = lr)

n0obcoder

import cv2
import torch.nn.functional as F

# Making a 'predict' function which would take the 'model' and the path of the 'test image' as inputs, and predict the class that the test image belongs to.
def predict(model, test_img_path):
    img = cv2.imread(test_img_path)

    # Visualizing the test image
    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))