Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| import torch.nn.functional as F | |
| from torchvision import datasets, transforms | |
| from torch.utils.data import DataLoader | |
| import streamlit as st | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| from PIL import Image | |
| # Define the Inception Module | |
| class InceptionModule(nn.Module): | |
| def __init__(self, in_channels): | |
| super(InceptionModule, self).__init__() | |
| self.branch1x1 = nn.Conv2d(in_channels, 64, kernel_size=1) | |
| self.branch3x3 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1) | |
| self.branch5x5 = nn.Conv2d(in_channels, 64, kernel_size=5, padding=2) | |
| self.branch_pool = nn.Conv2d(in_channels, 64, kernel_size=1) | |
| def forward(self, x): | |
| branch1x1 = self.branch1x1(x) | |
| branch3x3 = self.branch3x3(x) | |
| branch5x5 = self.branch5x5(x) | |
| branch_pool = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)(x) | |
| branch_pool = self.branch_pool(branch_pool) | |
| outputs = [branch1x1, branch3x3, branch5x5, branch_pool] | |
| return torch.cat(outputs, 1) | |
| # Define the Inception Network | |
| class InceptionNet(nn.Module): | |
| def __init__(self, num_classes=10): | |
| super(InceptionNet, self).__init__() | |
| self.inception1 = InceptionModule(in_channels=3) | |
| self.fc = nn.Linear(64 * 4 * 224 * 224, num_classes) | |
| def forward(self, x): | |
| x = self.inception1(x) | |
| x = x.view(x.size(0), -1) | |
| x = self.fc(x) | |
| return x | |
| # Training function | |
| def train(model, device, train_loader, optimizer, 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.cross_entropy(output, target) | |
| loss.backward() | |
| optimizer.step() | |
| if batch_idx % 10 == 0: | |
| print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}') | |
| # Define a simple transformation | |
| transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor() | |
| ]) | |
| # Initialize model and optimizer | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = InceptionNet().to(device) | |
| optimizer = optim.Adam(model.parameters(), lr=0.001) | |
| # Streamlit app | |
| st.title("InceptionNet Image Classifier") | |
| 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) | |
| image = transform(image).unsqueeze(0).to(device) | |
| st.write("Classifying...") | |
| model.eval() | |
| with torch.no_grad(): | |
| output = model(image) | |
| st.write("Output:", output.cpu().detach().numpy()) | |
| # Adjust Hyperparameters | |
| st.sidebar.title("Adjust Hyperparameters") | |
| learning_rate = st.sidebar.slider("Learning Rate", 0.0001, 0.01, 0.001) | |
| optimizer = optim.Adam(model.parameters(), lr=learning_rate) | |
| # Visualize Model's Predictions | |
| if st.sidebar.button("Train Model"): | |
| # Load dummy data for demonstration purposes | |
| train_loader = DataLoader( | |
| datasets.FakeData(transform=transform), | |
| batch_size=64, shuffle=True | |
| ) | |
| train(model, device, train_loader, optimizer, epoch=1) | |
| st.sidebar.write("Model Trained!") | |