[Python] How can I convert this code into binary classification?

The Actual Code link: https://github.com/mrdbourke/pytorc...rch_model_deployment_exercise_solutions.ipynb

I am trying to adjust the actual code above to work with my custom dataset for binary classification. I appreciate your help!

from google.colab import drive
drive.mount('/content/drive')

%%writefile get_data.py
import os
import zipfile

from pathlib import Path

import requests

# Setup path to data folder
data_path = Path("/content/drive/MyDrive/Dataset/")
image_path = data_path / "/content/drive/MyDrive/Dataset/Datas"

# If the image folder doesn't exist, download it and prepare it...
if image_path.is_dir():
print(f"{image_path} directory exists.")
else:
print(f"Did not find {image_path} directory, creating one...")
image_path.mkdir(parents=True, exist_ok=True)

%%writefile data_setup.py
"""
Contains functionality for creating PyTorch DataLoaders for
image classification data.
"""
import os

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

NUM_WORKERS = os.cpu_count()

def create_dataloaders(
train_dir: str,
test_dir: str,
transform: transforms.Compose,
batch_size: int,
num_workers: int=NUM_WORKERS
):

# Use ImageFolder to create dataset(s)
train_data = datasets.ImageFolder(train_dir, transform=transform)
test_data = datasets.ImageFolder(test_dir, transform=transform)

# Get class names
class_names = train_data.classes

# Turn images into data loaders
train_dataloader = DataLoader(
train_data,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True,
)
test_dataloader = DataLoader(
test_data,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=True,
)

return train_dataloader, test_dataloader, class_names

%%writefile engine.py
"""
Contains functions for training and testing a PyTorch model.
"""
import torch
from tqdm.auto import tqdm
from typing import Dict, List, Tuple

def train_step(model: torch.nn.Module,
dataloader: torch.utils.data.DataLoader,
loss_fn: torch.nn.Module,
optimizer: torch.optim.Optimizer,
device: torch.device) -> Tuple[float, float]:

# Put model in train mode
model.train()

# Setup train loss and train accuracy values
train_loss, train_acc = 0, 0

# Loop through data loader data batches
for batch, (X, y) in enumerate(dataloader):
# Send data to target device
X, y = X.to(device), y.to(device)
y = y.unsqueeze(1).float() # Ensure y has the shape [batch_size, 1]

# Print the shape of y for debugging
print(f"Train Step - y shape: {y.shape}")

# 1. Forward pass
y_pred = model(X)

# 2. Calculate and accumulate loss
loss = loss_fn(y_pred, y)
train_loss += loss.item()

# 3. Optimizer zero grad
optimizer.zero_grad()

# 4. Loss backward
loss.backward()

# 5. Optimizer step
optimizer.step()

# Calculate and accumulate accuracy metric across all batches
y_pred_class = torch.round(torch.sigmoid(y_pred))
train_acc += (y_pred_class == y).sum().item() / len(y_pred)

# Adjust metrics to get average loss and accuracy per batch
train_loss = train_loss / len(dataloader)
train_acc = train_acc / len(dataloader)
return train_loss, train_acc

def test_step(model: torch.nn.Module,
dataloader: torch.utils.data.DataLoader,
loss_fn: torch.nn.Module,
device: torch.device) -> Tuple[float, float]:

# Put model in eval mode
model.eval()

# Setup test loss and test accuracy values
test_loss, test_acc = 0, 0

# Turn on inference context manager
with torch.inference_mode():
# Loop through DataLoader batches
for batch, (X, y) in enumerate(dataloader):
# Send data to target device
X, y = X.to(device), y.to(device)
y = y.unsqueeze(1).float() # Ensure y has the shape [batch_size, 1]

# Print the shape of y for debugging
print(f"Test Step - y shape: {y.shape}")

# 1. Forward pass
test_pred_logits = model(X)

# 2. Calculate and accumulate loss
loss = loss_fn(test_pred_logits, y)
test_loss += loss.item()

# Calculate and accumulate accuracy
test_pred_labels = torch.round(torch.sigmoid(test_pred_logits))
test_acc += ((test_pred_labels == y).sum().item() / len(test_pred_labels))

# Adjust metrics to get average loss and accuracy per batch
test_loss = test_loss / len(dataloader)
test_acc = test_acc / len(dataloader)
return test_loss, test_acc

def train(model: torch.nn.Module,
train_dataloader: torch.utils.data.DataLoader,
test_dataloader: torch.utils.data.DataLoader,
optimizer: torch.optim.Optimizer,
loss_fn: torch.nn.Module,
epochs: int,
device: torch.device) -> Dict[str, List]:

# Create empty results dictionary
results = {"train_loss": [],
"train_acc": [],
"test_loss": [],
"test_acc": []
}

# Loop through training and testing steps for a number of epochs
for epoch in tqdm(range(epochs)):
train_loss, train_acc = train_step(model=model,
dataloader=train_dataloader,
loss_fn=loss_fn,
optimizer=optimizer,
device=device)
test_loss, test_acc = test_step(model=model,
dataloader=test_dataloader,
loss_fn=loss_fn,
device=device)

# Print out what's happening
print(
f"Epoch: {epoch+1} | "
f"train_loss: {train_loss:.4f} | "
f"train_acc: {train_acc:.4f} | "
f"test_loss: {test_loss:.4f} | "
f"test_acc: {test_acc:.4f}"
)

# Update results dictionary
results["train_loss"].append(train_loss)
results["train_acc"].append(train_acc)
results["test_loss"].append(test_loss)
results["test_acc"].append(test_acc)

# Return the filled results at the end of the epochs
return results

%%writefile model_builder.py
"""
Contains PyTorch model code to instantiate a TinyVGG model.
"""
import torch
from torch import nn

class TinyVGG(nn.Module):

def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
super().__init__()
self.conv_block_1 = nn.Sequential(
nn.Conv2d(in_channels=input_shape,
out_channels=hidden_units,
kernel_size=3,
stride=1,
padding=0),
nn.ReLU(),
nn.Conv2d(in_channels=hidden_units,
out_channels=hidden_units,
kernel_size=3,
stride=1,
padding=0),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2,
stride=2)
)
self.conv_block_2 = nn.Sequential(
nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
nn.ReLU(),
nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=0),
nn.ReLU(),
nn.MaxPool2d(2)
)
self.classifier = nn.Sequential(
nn.Flatten(),
# Where did this in_features shape come from?
# It's because each layer of our network compresses and changes the shape of our inputs data.
nn.Linear(in_features=hidden_units*13*13,
out_features=1)
)

def forward(self, x: torch.Tensor):
x = self.conv_block_1(x)
x = self.conv_block_2(x)
x = self.classifier(x)
return x
# return self.classifier(self.block_2(self.block_1(x))) # <- leverage the benefits of operator fusion

%%writefile utils.py
"""
Contains various utility functions for PyTorch model training and saving.
"""
import torch
from pathlib import Path

def save_model(model: torch.nn.Module,
target_dir: str,
model_name: str):

# Create target directory
target_dir_path = Path(target_dir)
target_dir_path.mkdir(parents=True,
exist_ok=True)

# Create model save path
assert model_name.endswith(".pth") or model_name.endswith(".pt"), "model_name should end with '.pt' or '.pth'"
model_save_path = target_dir_path / model_name

# Save the model state_dict()
print(f"[INFO] Saving model to: {model_save_path}")
torch.save(obj=model.state_dict(),
f=model_save_path)
%%writefile train.py
"""
Trains a PyTorch image classification model using device-agnostic code.
"""

import os
import argparse

import torch

from torchvision import transforms

import data_setup, engine, model_builder, utils

# Create a parser
parser = argparse.ArgumentParser(description="Get some hyperparameters.")

# Get an arg for num_epochs
parser.add_argument("--num_epochs",
default=10,
type=int,
help="the number of epochs to train for")

# Get an arg for batch_size
parser.add_argument("--batch_size",
default=32,
type=int,
help="number of samples per batch")

# Get an arg for hidden_units
parser.add_argument("--hidden_units",
default=10,
type=int,
help="number of hidden units in hidden layers")

# Get an arg for learning_rate
parser.add_argument("--learning_rate",
default=0.001,
type=float,
help="learning rate to use for model")

# Create an arg for training directory
parser.add_argument("--train_dir",
default="/content/drive/MyDrive/Dataset/Datas/Train",
type=str,
help="directory file path to training data in standard image classification format")

# Create an arg for test directory
parser.add_argument("--test_dir",
default="/content/drive/MyDrive/Dataset/Datas/Test",
type=str,
help="directory file path to testing data in standard image classification format")

# Get our arguments from the parser
args = parser.parse_args()

# Setup hyperparameters
NUM_EPOCHS = args.num_epochs
BATCH_SIZE = args.batch_size
HIDDEN_UNITS = args.hidden_units
LEARNING_RATE = args.learning_rate
print(f"[INFO] Training a model for {NUM_EPOCHS} epochs with batch size {BATCH_SIZE} using {HIDDEN_UNITS} hidden units and a learning rate of {LEARNING_RATE}")

# Setup directories
train_dir = args.train_dir
test_dir = args.test_dir
print(f"[INFO] Training data file: {train_dir}")
print(f"[INFO] Testing data file: {test_dir}")

# Setup target device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Create transforms
data_transform = transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor()
])

# Create DataLoaders with help from data_setup.py
train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
train_dir=train_dir,
test_dir=test_dir,
transform=data_transform,
batch_size=BATCH_SIZE
)

# Create model with help from model_builder.py
model = model_builder.TinyVGG(
input_shape=3,
hidden_units=HIDDEN_UNITS,
output_shape=1
).to(device)

# Set loss and optimizer
loss_fn = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(),
lr=LEARNING_RATE)

# Start training with help from engine.py
engine.train(model=model,
train_dataloader=train_dataloader,
test_dataloader=test_dataloader,
loss_fn=loss_fn,
optimizer=optimizer,
epochs=NUM_EPOCHS,
device=device)

# Save the model with help from utils.py
utils.save_model(model=model,
target_dir="models",
model_name="05_going_modular_script_mode_tinyvgg_model.pth")

!python train.py --num_epochs 10 --batch_size 32 --hidden_units 128 --learning_rate 0.001



The Actual Code link: https://github.com/mrdbourke/pytorc...rch_model_deployment_exercise_solutions.ipynb

I am trying to adjust the actual code above to work with my custom dataset for binary classification. I appreciate your help!

I have 512x512 RGB images, but the above code is resizing them to 64x64. I am using BCEWithLogitsLoss, and I am trying to use unsqueeze, but I can't transition from multi-class classification to binary classification. I have two classes: diseased and healthy, and I have a custom dataset. I keep getting the errors: "Target size (torch.Size([32])) must be the same as input size (torch.Size([32, 1]))" or "RuntimeError: result type Float can't be cast to the desired output type Long."

What part of the code is faulty, and how should I proceed? I would really appreciate your help!

Continue reading...