Spaces:
Running
Running
| # Install required packages | |
| # !pip install streamlit torch torchvision matplotlib | |
| # Import Libraries | |
| import streamlit as st | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| import torchvision # Add this import | |
| from torchvision import datasets, models, transforms | |
| from torch.utils.data import DataLoader, Subset | |
| import numpy as np | |
| import time | |
| import copy # Add this import | |
| import matplotlib.pyplot as plt | |
| # Streamlit Interface | |
| st.title("Simple ResNet Fine-Tuning Example") | |
| # User Inputs | |
| st.sidebar.header("Model Parameters") | |
| batch_size = st.sidebar.number_input("Batch Size", value=32) | |
| num_epochs = st.sidebar.number_input("Number of Epochs", value=5) | |
| learning_rate = st.sidebar.number_input("Learning Rate", value=0.001) | |
| # Data Preparation Section | |
| st.markdown(""" | |
| ### Data Preparation | |
| We will use a small subset of the CIFAR-10 dataset for quick experimentation. The dataset will be split into training and validation sets, and transformations will be applied to normalize the data. | |
| """) | |
| transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) | |
| ]) | |
| train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) | |
| val_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) | |
| # Using only 1000 samples for simplicity | |
| subset_indices = list(range(1000)) | |
| train_size = int(0.8 * len(subset_indices)) | |
| val_size = len(subset_indices) - train_size | |
| train_indices = subset_indices[:train_size] | |
| val_indices = subset_indices[train_size:] | |
| train_dataset = Subset(train_dataset, train_indices) | |
| val_dataset = Subset(val_dataset, val_indices) | |
| train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4) | |
| val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4) | |
| dataloaders = {'train': train_loader, 'val': val_loader} | |
| class_names = datasets.CIFAR10(root='./data', download=False).classes | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| # 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) | |
| fig, ax = plt.subplots() | |
| ax.imshow(inp) | |
| if title is not None: | |
| ax.set_title(title) | |
| st.pyplot(fig) | |
| inputs, classes = next(iter(dataloaders['train'])) | |
| out = torchvision.utils.make_grid(inputs) | |
| 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 dataset. | |
| """) | |
| # Load Pre-trained ResNet Model | |
| model_ft = models.resnet | |