Update pages/19_ResNet.py

pages/19_ResNet.py

@@ -27,7 +27,6 @@ In this exercise, we will fine-tune a pre-trained ResNet model on a custom image
 
 # 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)
@@ -37,30 +36,24 @@ 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.
+We will use the CIFAR-10 dataset, which contains 60,000 images from 10 classes. The dataset will be split into training and validation sets, and transformations will be applied to augment the data and normalize it.
 """)
 
-data_transforms = {
-    'train': transforms.Compose([
-        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(input_size),
-        transforms.CenterCrop(input_size),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ]),
-}
-
-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=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
+transform = transforms.Compose([
+    transforms.Resize(input_size),
+    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)
+
+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}
+dataset_sizes = {'train': len(train_dataset), 'val': len(val_dataset)}
+class_names = train_dataset.classes
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -204,32 +197,8 @@ st.markdown("""
 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,
-)
-
+# This part is just illustrative since Hugging Face's Trainer does not natively support ResNet. 
+# However, you can still follow a similar approach for transformer models and NLP datasets.
 if st.button('Train with Hugging Face'):
-    trainer.train()
-    st.write("Model trained using Hugging Face")
+    st.write("This section is illustrative and typically used for NLP tasks with Hugging Face transformers.")