Update pages/19_ResNet.py

pages/19_ResNet.py

@@ -1,8 +1,8 @@
-# Install necessary packages
-# Ensure you have PyTorch, torchvision, and Streamlit installed
-# You can install them using pip if you haven't already:
-# pip install torch torchvision streamlit
+# Install required packages
+# !pip install streamlit torch torchvision matplotlib datasets transformers
 
+# Import Libraries
+import streamlit as st
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -12,64 +12,118 @@ import numpy as np
 import time
 import os
 import copy
-import streamlit as st
-from PIL import Image
 import matplotlib.pyplot as plt
-import torchvision.transforms as T
+from transformers import Trainer, TrainingArguments
+from datasets import load_dataset
+
+# Streamlit Interface
+st.title("Fine-Tuning ResNet for Custom Image Classification")
+
+# Introduction Section
+st.markdown("""
+### Introduction
+In this exercise, we will fine-tune a pre-trained ResNet model on a custom image classification task using PyTorch. The ResNet (Residual Network) architecture helps in training very deep neural networks by using skip connections to mitigate the vanishing gradient problem.
+""")
+
+# User Inputs
+st.sidebar.header("Model Parameters")
+data_dir = st.sidebar.text_input("Path to Dataset Directory", 'path_to_caltech101_dataset')
+input_size = st.sidebar.number_input("Input Size", value=224)
+batch_size = st.sidebar.number_input("Batch Size", value=32)
+num_epochs = st.sidebar.number_input("Number of Epochs", value=25)
+learning_rate = st.sidebar.number_input("Learning Rate", value=0.001)
+momentum = st.sidebar.number_input("Momentum", value=0.9)
+
+# Data Preparation Section
+st.markdown("""
+### Data Preparation
+We will use the Caltech 101 dataset, which contains images from 101 object categories. The dataset will be split into training and validation sets, and transformations will be applied to augment the data and normalize it.
+""")
 
-# Data transformations
 data_transforms = {
     'train': transforms.Compose([
-        transforms.RandomResizedCrop(224),
+        transforms.RandomResizedCrop(input_size),
         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.Resize(input_size),
+        transforms.CenterCrop(input_size),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
     ]),
 }
 
-# Load datasets
-data_dir = 'path/to/data'
 image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
                   for x in ['train', 'val']}
-dataloaders = {x: DataLoader(image_datasets[x], batch_size=32, shuffle=True, num_workers=4)
+dataloaders = {x: DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4)
                for x in ['train', 'val']}
 dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
 class_names = image_datasets['train'].classes
+
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# Load the pre-trained model
+# Visualize a few training images
+st.markdown("#### Sample Training Images")
+def imshow(inp, title=None):
+    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)
+    plt.imshow(inp)
+    if title is not None:
+        plt.title(title)
+    plt.pause(0.001)
+
+inputs, classes = next(iter(dataloaders['train']))
+out = torchvision.utils.make_grid(inputs)
+st.pyplot(imshow(out, title=[class_names[x] for x in classes]))
+
+# Model Preparation Section
+st.markdown("""
+### Model Preparation
+We will use a pre-trained ResNet-18 model and fine-tune the final fully connected layer to match the number of classes in our custom dataset.
+""")
+
+# Load Pre-trained ResNet Model
 model_ft = models.resnet18(pretrained=True)
 num_ftrs = model_ft.fc.in_features
 model_ft.fc = nn.Linear(num_ftrs, len(class_names))
+
 model_ft = model_ft.to(device)
 
-# Define loss function and optimizer
+# Define Loss Function and Optimizer
 criterion = nn.CrossEntropyLoss()
-optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
-scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
+optimizer_ft = optim.SGD(model_ft.parameters(), lr=learning_rate, momentum=momentum)
+exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
+
+# Training Section
+st.markdown("""
+### Training
+We will train the model using stochastic gradient descent (SGD) with momentum and a learning rate scheduler. The training and validation loss and accuracy will be plotted to monitor the training process.
+""")
 
-# Training and evaluation functions
+# Train and Evaluate the Model
 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
+    train_loss_history = []
+    val_loss_history = []
+    train_acc_history = []
+    val_acc_history = []
 
     for epoch in range(num_epochs):
-        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
-        print('-' * 10)
+        st.write('Epoch {}/{}'.format(epoch, num_epochs - 1))
+        st.write('-' * 10)
 
         for phase in ['train', 'val']:
             if phase == 'train':
-                model.train()  
+                model.train()
             else:
-                model.eval()   
+                model.eval()
 
             running_loss = 0.0
             running_corrects = 0
@@ -98,66 +152,84 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
             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))
+            st.write('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
+
+            if phase == 'train':
+                train_loss_history.append(epoch_loss)
+                train_acc_history.append(epoch_acc)
+            else:
+                val_loss_history.append(epoch_loss)
+                val_acc_history.append(epoch_acc)
 
             if phase == 'val' and epoch_acc > best_acc:
                 best_acc = epoch_acc
                 best_model_wts = copy.deepcopy(model.state_dict())
 
-        print()
+        st.write()
 
     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))
+    st.write('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
+    st.write('Best val Acc: {:4f}'.format(best_acc))
 
     model.load_state_dict(best_model_wts)
+    
+    # Plot training history
+    epochs_range = range(num_epochs)
+    plt.figure(figsize=(10, 5))
+    plt.subplot(1, 2, 1)
+    plt.plot(epochs_range, train_loss_history, label='Training Loss')
+    plt.plot(epochs_range, val_loss_history, label='Validation Loss')
+    plt.legend(loc='upper right')
+    plt.title('Training and Validation Loss')
+
+    plt.subplot(1, 2, 2)
+    plt.plot(epochs_range, train_acc_history, label='Training Accuracy')
+    plt.plot(epochs_range, val_acc_history, label='Validation Accuracy')
+    plt.legend(loc='lower right')
+    plt.title('Training and Validation Accuracy')
+    plt.show()
+    st.pyplot(plt)
+    
     return model
 
-# Train the model
-model_ft = train_model(model_ft, criterion, optimizer_ft, scheduler, num_epochs=25)
-
-# Save the trained model
-torch.save(model_ft.state_dict(), 'model_ft.pth')
+if st.button('Train Model'):
+    model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs)
+    # Save the Model
+    torch.save(model_ft.state_dict(), 'fine_tuned_resnet.pth')
+    st.write("Model saved as 'fine_tuned_resnet.pth'")
+
+# Hugging Face Integration Section
+st.markdown("""
+### Hugging Face Integration
+We will use the Hugging Face library to load the dataset and prepare it for training. This integration will allow us to leverage the benefits of Hugging Face's powerful tools and APIs.
+""")
+
+dataset = load_dataset('caltech101', split='train')
+
+def preprocess_function(examples):
+    return {'pixel_values': [data_transforms['train'](image) for image in examples['image']], 'labels': examples['label']}
+
+dataset = dataset.map(preprocess_function, batched=True)
+
+training_args = TrainingArguments(
+    output_dir='./results',
+    evaluation_strategy="epoch",
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=8,
+    num_train_epochs=3,
+    save_strategy="epoch",
+    logging_dir='./logs',
+)
+
+trainer = Trainer(
+    model=model_ft,
+    args=training_args,
+    train_dataset=dataset['train'],
+    eval_dataset=dataset['val'],
+    tokenizer=None,
+)
+
+if st.button('Train with Hugging Face'):
+    trainer.train()
+    st.write("Model trained using Hugging Face")
 
-# Streamlit Interface
-st.title("Image Classification with Fine-tuned ResNet")
-
-uploaded_file = st.file_uploader("Choose an image...", type="jpg")
-
-if uploaded_file is not None:
-    image = Image.open(uploaded_file)
-    st.image(image, caption='Uploaded Image.', use_column_width=True)
-    st.write("")
-    st.write("Classifying...")
-
-    model_ft = models.resnet18(pretrained=True)
-    num_ftrs = model_ft.fc.in_features
-    model_ft.fc = nn.Linear(num_ftrs, len(class_names))
-    model_ft.load_state_dict(torch.load('model_ft.pth'))
-    model_ft = model_ft.to(device)
-    model_ft.eval()
-
-    preprocess = T.Compose([
-        T.Resize(256),
-        T.CenterCrop(224),
-        T.ToTensor(),
-        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-
-    img = preprocess(image).unsqueeze(0)
-    img = img.to(device)
-
-    with torch.no_grad():
-        outputs = model_ft(img)
-        _, preds = torch.max(outputs, 1)
-        predicted_class = class_names[preds[0]]
-
-    st.write(f"Predicted Class: {predicted_class}")
-
-    # Plotting the image with matplotlib
-    fig, ax = plt.subplots()
-    ax.imshow(image)
-    ax.set_title(f"Predicted: {predicted_class}")
-    st.pyplot(fig)