Build an AI-Generated Image Detector

Fine-tune ResNet50 to classify real vs. AI-generated images with PyTorch.

Introduction

30 Second Summary

AI-generated images are everywhere. From social media feeds to news articles, it is getting harder to tell what is real and what a machine created.

In this project, you will train a deep learning model that can look at any image and tell you whether it is a real photograph or generated by AI. You will use the same transfer learning workflow that powers production machine learning systems.

What You'll Build

You will fine-tune a pre-trained ResNet50 model on 120,000 images using PyTorch in Google Colab to build a binary classifier that detects AI-generated content.

By the end of this project, you'll have:

  • A trained ResNet50 model achieving 95%+ validation accuracy at classifying images as real or AI-generated.
  • A complete training pipeline in Google Colab covering data loading, transfer learning, and model evaluation.
  • Visualizations including training loss curves and a confusion matrix that prove your model learned effectively.
  • An inference function that classifies any new image you feed it as real or fake with a confidence score.
  • Secret Mission: Add Grad-CAM visualization to see what regions of an image the model focuses on when making its prediction.

Are there any prerequisites?

You should have some Python programming experience and a basic understanding of what neural networks are. You do not need prior experience with PyTorch or image classification.

Before We Start

Before We Start...

Set Up Your Environment

To build an AI image detector, you need a computing environment with GPU acceleration and access to a large dataset.

In this step, you will set up Google Colab with a GPU runtime, configure the Kaggle API to download your training data, and confirm that PyTorch can access the GPU.

In this step, get ready to:
  • Open Google Colab and connect to a T4 GPU runtime.
  • Configure the Kaggle API and download the CIFAKE dataset.
  • Verify PyTorch is working and can access the GPU.
Open Google Colab and connect to a GPU

Google Colab gives you free access to GPUs in your browser. Training a deep learning model without a GPU would take hours instead of minutes.

  • Open colab.research.google.com in your browser and sign in with your Google account.
  • Click New notebook to create a blank notebook.
  • In the top menu, click Runtime.
  • Select Change runtime type.
  • In the Hardware accelerator dropdown, select GPU.
  • Click Save.

Why do we need a GPU?

Deep learning models process thousands of images during training. A GPU handles these calculations in parallel, making training 10-50x faster than a CPU.

Colab's free tier provides access to NVIDIA T4 GPUs, which are more than enough for this project.

Don't see a GPU option?

Make sure you are signed in to your Google account. GPU access on the free tier is not guaranteed and depends on availability. If no GPU is available, try again in a few minutes or continue with CPU (training will be slower but still works).

Help me connect to a GPU in Google Colab.

Configure Kaggle and download the CIFAKE dataset

The CIFAKE dataset contains 120,000 images split evenly between real photos and AI-generated fakes. You will download it using the Kaggle API, which requires a free Kaggle account and an API token.

  • Go to kaggle.com and sign in (or create a free account if you do not have one).
  • Click your profile icon in the top right, then select Settings.
  • Scroll down to the API section and click Create New API Token.
  • This downloads a file called kaggle.json to your computer.

What is kaggle.json?

This file contains your Kaggle username and API key. It lets you download datasets from Kaggle programmatically without logging in through the browser each time.

Now you need to upload kaggle.json to your Colab notebook and write the code that downloads the dataset.

  • In your Colab notebook, click the folder icon on the left sidebar to open the Files panel.
  • Click the Upload button (the file icon with an up arrow) at the top of the Files panel.
  • Select the kaggle.json file you just downloaded.
  • Wait for the upload to complete. The file appears under /content/kaggle.json in the file browser.

Before running the download code, you need to set up your imports. These libraries will be used throughout the entire notebook.

  • Add the following code to the first cell of your notebook and run it:
# AI-Generated Image Detector - Google Colab Notebook
# Fine-tunes ResNet50 on CIFAKE dataset to classify real vs AI-generated images

# ============================================================
# Cell 1: Setup and Imports
# ============================================================
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
from torchvision import datasets, transforms
from torchvision.models import resnet50, ResNet50_Weights
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

# Verify GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

What does this code do?

  • os handles file system operations like moving your Kaggle credentials into place.
  • torch is the PyTorch deep learning framework you will use to build and train the model.
  • torchvision provides pre-trained models (ResNet50), image transforms, and dataset loaders.
  • matplotlib and numpy are for plotting training curves and handling image arrays.
  • sklearn.metrics provides the confusion matrix you will use to evaluate model performance.
  • The last two lines check whether a GPU is available and store the result in device for later use.

You should see Using device: cuda printed below the cell. This confirms your notebook has GPU access.

  • Now add a second cell below with the dataset download code and run it:
# ============================================================
# Cell 2: Download CIFAKE Dataset
# ============================================================
# Upload kaggle.json first (use Colab file upload)
os.makedirs("/root/.kaggle", exist_ok=True)
os.rename("/content/kaggle.json", "/root/.kaggle/kaggle.json")
os.chmod("/root/.kaggle/kaggle.json", 0o600)

os.system("kaggle datasets download -d birdy654/cifake-real-and-ai-generated-synthetic-images -p /content/data")
os.system("unzip -q /content/data/cifake-real-and-ai-generated-synthetic-images.zip -d /content/data/cifake")

What does this code do?

  • os.makedirs("/root/.kaggle", exist_ok=True) creates the directory where Kaggle expects to find your credentials.
  • os.rename moves your uploaded kaggle.json from the upload location to the Kaggle config directory.
  • os.chmod restricts file permissions so only you can read the credentials (Kaggle requires this).
  • The os.system calls download the CIFAKE dataset (about 600 MB) and extract it to /content/data/cifake.

The download takes 1-2 minutes depending on connection speed. You should see progress output from the Kaggle CLI.

Seeing a 404 or permission error?

  • Make sure you uploaded kaggle.json before running the cell. If you ran Cell 2 first, re-upload kaggle.json (the rename moved or failed on the original) and re-run.
  • Check that your Kaggle account has accepted the dataset's terms. Visit the CIFAKE dataset page on Kaggle and click Download once manually to accept any rules.
  • If you see a credentials error, regenerate your API token on Kaggle and re-upload the new kaggle.json.

Help me troubleshoot the Kaggle dataset download in Colab.

Verify GPU availability

When you ran Cell 1, the last line printed the device your notebook is using. This is your confirmation that the environment is ready.

  • Scroll up to the output of Cell 1 and confirm it says Using device: cuda.
  • In the Files panel on the left, expand /content/data/cifake and confirm you see train and test folders.

What if it says cpu instead of cuda?

If your output shows Using device: cpu, your notebook did not connect to a GPU. Go back to Runtime > Change runtime type and make sure GPU is selected.

If GPU is selected but still shows cpu, Colab may be out of free GPU capacity. You can still complete this project on CPU. Training will take longer (roughly 10-15 minutes per epoch instead of 1-2 minutes).

✔️ Awesome, I've got everything!

Great. Make sure both cells ran without errors and you can see the CIFAKE dataset folders in your file browser.

ⓧ I'd like to double check the full code

Here is the complete code you should have across both cells in your notebook so far:

# AI-Generated Image Detector - Google Colab Notebook
# Fine-tunes ResNet50 on CIFAKE dataset to classify real vs AI-generated images

# ============================================================
# Cell 1: Setup and Imports
# ============================================================
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
from torchvision import datasets, transforms
from torchvision.models import resnet50, ResNet50_Weights
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

# Verify GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ============================================================
# Cell 2: Download CIFAKE Dataset
# ============================================================
# Upload kaggle.json first (use Colab file upload)
os.makedirs("/root/.kaggle", exist_ok=True)
os.rename("/content/kaggle.json", "/root/.kaggle/kaggle.json")
os.chmod("/root/.kaggle/kaggle.json", 0o600)

os.system("kaggle datasets download -d birdy654/cifake-real-and-ai-generated-synthetic-images -p /content/data")
os.system("unzip -q /content/data/cifake-real-and-ai-generated-synthetic-images.zip -d /content/data/cifake")

Your environment is ready with GPU access and 120,000 training images. Next up, you will load this dataset into PyTorch and explore what real vs. AI-generated images look like.

Load and Explore the Dataset

Your environment is ready and the CIFAKE dataset is sitting in your Colab filesystem. Now you need to get that raw image data into a format PyTorch can actually work with.

Neural networks expect images as tensors with specific dimensions and normalized pixel values. In this step, you will define the preprocessing pipeline that matches what ResNet50 expects, load the images into datasets, and visualize samples from both classes.

In this step, get ready to:
  • Define image transforms that match ResNet50's expected input format.
  • Load the CIFAKE dataset and create DataLoaders for batched training.
  • Visualize sample images from the REAL and FAKE classes.
Define the preprocessing transforms

ResNet50 was trained on ImageNet images preprocessed in a specific way. For the model to make sense of new images, you need to apply the same preprocessing steps it was trained with.

  • In your Colab notebook, create a new code cell below your dataset download cell.
  • Define the training and test transforms by pasting this code:
# ============================================================
# Cell 3: Load and Explore Dataset
# ============================================================
# Define transforms matching ResNet50 IMAGENET1K_V2 preprocessing
train_transform = transforms.Compose([
    transforms.Resize(232),
    transforms.CenterCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

test_transform = transforms.Compose([
    transforms.Resize(232),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

What does this code do?

Each transform step prepares the image for ResNet50's input expectations:

  • Resize(232) scales the image to 232 pixels on the shortest side.
  • CenterCrop(224) takes a 224x224 center crop, which is the exact resolution ResNet50 expects.
  • RandomHorizontalFlip() is a data augmentation technique that randomly mirrors images during training. This helps the model generalize instead of memorizing specific orientations.
  • Normalize adjusts pixel values to match the ImageNet distribution the model was originally trained on.

The test transform skips RandomHorizontalFlip because you want consistent, deterministic results during evaluation.

You won't see a change yet. The next chunk uses these transforms to load images.

Load the dataset and create DataLoaders

PyTorch's ImageFolder class automatically loads images from a directory structure where each subfolder is a class label. The CIFAKE dataset already uses this exact structure with REAL/ and FAKE/ subfolders.

  • Add this code below your transforms in the same cell:
# Load datasets using ImageFolder (expects subdirectories per class)
train_dataset = datasets.ImageFolder("/content/data/cifake/train", transform=train_transform)
test_dataset = datasets.ImageFolder("/content/data/cifake/test", transform=test_transform)

print(f"Training samples: {len(train_dataset)}")
print(f"Test samples: {len(test_dataset)}")
print(f"Classes: {train_dataset.classes}")

# Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=2)

What does this code do?

ImageFolder scans the folder structure and assigns class labels alphabetically. Since FAKE comes before REAL alphabetically, the classes list will be ['FAKE', 'REAL'] with indices 0 and 1.

DataLoader wraps the dataset to serve images in batches of 32. Shuffling the training data prevents the model from learning the order of images instead of their content. The test loader does not shuffle so results stay reproducible.

  • Run the cell by pressing Shift+Enter.
  • You should see output showing Training samples: 100000, Test samples: 20000, and Classes: ['FAKE', 'REAL'].

Seeing a FileNotFoundError?

  • Make sure the dataset path is exactly /content/data/cifake/train. If you used a different extraction path in Step 1, adjust accordingly.
  • Confirm the unzip completed successfully in your previous cell. Re-run the download cell if needed.

Help me debug a FileNotFoundError when loading my CIFAKE dataset with ImageFolder in Colab.

Visualize sample images from both classes

Before training, it helps to see what the model will learn from. You will write a helper function that undoes the normalization and displays a grid of images with their class labels.

  • Add this visualization code below in the same cell:
# Visualize sample images
def show_samples(dataset, num_images=8):
    fig, axes = plt.subplots(2, num_images // 2, figsize=(14, 6))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    for i, ax in enumerate(axes.flat):
        img, label = dataset[i * (len(dataset) // num_images)]
        img = img.numpy().transpose(1, 2, 0)
        img = std * img + mean
        img = np.clip(img, 0, 1)
        ax.imshow(img)
        ax.set_title(dataset.classes[label])
        ax.axis("off")
    plt.tight_layout()
    plt.show()

show_samples(train_dataset)

What does this code do?

  • The function samples 8 images evenly spaced across the dataset so you see variety.
  • img.numpy().transpose(1, 2, 0) converts the tensor from PyTorch's channel-first format (C, H, W) to matplotlib's channel-last format (H, W, C).
  • std * img + mean reverses the normalization so colors display correctly.
  • np.clip(img, 0, 1) clamps any pixel values that drifted outside the valid 0-1 range during denormalization.
  • Run the cell by pressing Shift+Enter.
  • You should see a 2x4 grid of images. Each image is labeled either REAL or FAKE above it.

Images look washed out or weirdly colored?

  • Check that the mean and std arrays inside show_samples match the ones in your transforms. Both must use [0.485, 0.456, 0.406] and [0.229, 0.224, 0.225].
  • If images appear blank, make sure train_dataset loaded successfully in the previous chunk (100,000 samples).

Help me fix washed out or incorrectly colored images in my matplotlib visualization.

What should I notice in the images?

The REAL images come from CIFAR-10 and appear slightly pixelated because they are originally 32x32, upscaled to 224x224. The FAKE images were generated by Stable Diffusion to mimic CIFAR-10 classes.

Both sets look similar at first glance. That is exactly why you need a neural network to tell them apart.

✔️ Awesome, I've got everything!

Great. Double check you have saved all cells in your notebook.

ⓧ I'd like to double check the full code

Here is the full notebook code up to this point:

# AI-Generated Image Detector - Google Colab Notebook
# Fine-tunes ResNet50 on CIFAKE dataset to classify real vs AI-generated images

# ============================================================
# Cell 1: Setup and Imports
# ============================================================
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
from torchvision import datasets, transforms
from torchvision.models import resnet50, ResNet50_Weights
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

# Verify GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ============================================================
# Cell 2: Download CIFAKE Dataset
# ============================================================
# Upload kaggle.json first (use Colab file upload)
os.makedirs("/root/.kaggle", exist_ok=True)
os.rename("/content/kaggle.json", "/root/.kaggle/kaggle.json")
os.chmod("/root/.kaggle/kaggle.json", 0o600)

os.system("kaggle datasets download -d birdy654/cifake-real-and-ai-generated-synthetic-images -p /content/data")
os.system("unzip -q /content/data/cifake-real-and-ai-generated-synthetic-images.zip -d /content/data/cifake")

# ============================================================
# Cell 3: Load and Explore Dataset
# ============================================================
# Define transforms matching ResNet50 IMAGENET1K_V2 preprocessing
train_transform = transforms.Compose([
    transforms.Resize(232),
    transforms.CenterCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

test_transform = transforms.Compose([
    transforms.Resize(232),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Load datasets using ImageFolder (expects subdirectories per class)
train_dataset = datasets.ImageFolder("/content/data/cifake/train", transform=train_transform)
test_dataset = datasets.ImageFolder("/content/data/cifake/test", transform=test_transform)

print(f"Training samples: {len(train_dataset)}")
print(f"Test samples: {len(test_dataset)}")
print(f"Classes: {train_dataset.classes}")

# Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=2)

# Visualize sample images
def show_samples(dataset, num_images=8):
    fig, axes = plt.subplots(2, num_images // 2, figsize=(14, 6))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    for i, ax in enumerate(axes.flat):
        img, label = dataset[i * (len(dataset) // num_images)]
        img = img.numpy().transpose(1, 2, 0)
        img = std * img + mean
        img = np.clip(img, 0, 1)
        ax.imshow(img)
        ax.set_title(dataset.classes[label])
        ax.axis("off")
    plt.tight_layout()
    plt.show()

show_samples(train_dataset)

Your data pipeline is ready. Next up, you will build the actual model by loading a pre-trained ResNet50 and adapting it for binary classification.

Build the Model with Transfer Learning

You have 120,000 images loaded and ready to go. Now you need a model that can look at any image and decide whether it is real or AI-generated.

Training a deep learning model from scratch would take millions of images and days of compute. Instead, you will use transfer learning. You will take a model that already knows how to "see" (trained on millions of photos) and teach it one new skill: spotting AI-generated images.

In this step, get ready to:
  • Load a pre-trained ResNet50 model and freeze its learned features.
  • Replace the model's output layer for binary classification (REAL vs FAKE).
  • Define the loss function and optimizer for training.
Load pre-trained ResNet50 and freeze backbone layers

ResNet50 is a 50-layer neural network that was trained on over a million images from the ImageNet dataset. It already knows how to detect edges, textures, shapes, and objects. You will keep all of that knowledge frozen and only teach it one new thing: the difference between real and AI-generated images.

  • Create a new code cell in your notebook and add the following code to load the model and freeze its parameters:
# Load pre-trained ResNet50
model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

# Freeze all backbone layers
for param in model.parameters():
    param.requires_grad = False

What does this code do?

The first line downloads ResNet50 with its pre-trained weights from the IMAGENET1K_V2 checkpoint. These weights encode everything the model learned from 1.2 million ImageNet photos.

The for loop sets requires_grad = False on every parameter. This "freezes" the backbone so its weights will not change during training. Only the new layer you add next will learn.

Why freeze layers instead of training everything?

ResNet50 has over 25 million parameters. Training all of them on your dataset would take hours and risk overwriting the valuable features it already knows.

By freezing the backbone, you get a model that trains in minutes and only needs to learn the final classification decision. This is how most production ML teams apply deep learning to new tasks.

Replace the classification head for binary output

ResNet50 was originally trained to classify images into 1,000 categories (dog breeds, car types, food). Your task only has 2 categories: REAL and FAKE. You need to swap out the final layer so the model outputs a prediction for your two classes.

Define the loss function and optimizer

The loss function measures how wrong the model's predictions are. The optimizer decides how to update the model's weights to reduce that loss. You will use CrossEntropyLoss (standard for classification) and the Adam optimizer targeting only the new FC layer.

  • Add the following code below the freeze loop in the same cell to complete your model setup:
# Replace the final FC layer for binary classification
model.fc = nn.Linear(2048, 2)

# Move model to GPU
model = model.to(device)

# Define loss and optimizer (only train the new FC layer)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

print(f"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")

What does this code do?

  • model.fc = nn.Linear(2048, 2) replaces ResNet50's original 1000-class output layer with a new 2-class layer. The 2048 is the number of features the backbone outputs. The 2 matches your two classes (REAL and FAKE).
  • model.to(device) moves the entire model to your GPU for faster computation.
  • CrossEntropyLoss is the standard loss function for classification tasks. It penalizes confident wrong predictions more heavily.
  • model.fc.parameters() tells the optimizer to only update the new FC layer's weights. The frozen backbone stays untouched.
  • Run the cell by pressing **Shift+Enter**.
  • You should see the output Trainable parameters: 4098 printed below the cell.

Why only 4,098 trainable parameters?

Your new nn.Linear(2048, 2) layer has a 2048x2 weight matrix (4,096 values) plus 2 bias values. That is 4,098 total. Out of ResNet50's 25 million parameters, you are only training these 4,098. Everything else stays frozen.

Not seeing the expected output?

  • If you see Trainable parameters: 0, make sure the model.fc = nn.Linear(2048, 2) line comes AFTER the freeze loop. New layers default to requires_grad = True, but only if created after the freeze.
  • If you see a large number like 25 million, the freeze loop did not run. Check that the for param loop is indented correctly and sets param.requires_grad = False.

Why is my trainable parameter count wrong?

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook before moving on.

ⓧ I'd like to double check the full code

Here is the complete notebook up to this point:

# AI-Generated Image Detector - Google Colab Notebook
# Fine-tunes ResNet50 on CIFAKE dataset to classify real vs AI-generated images

# ============================================================
# Cell 1: Setup and Imports
# ============================================================
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
from torchvision import datasets, transforms
from torchvision.models import resnet50, ResNet50_Weights
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

# Verify GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ============================================================
# Cell 2: Download CIFAKE Dataset
# ============================================================
os.makedirs("/root/.kaggle", exist_ok=True)
os.rename("/content/kaggle.json", "/root/.kaggle/kaggle.json")
os.chmod("/root/.kaggle/kaggle.json", 0o600)

os.system("kaggle datasets download -d birdy654/cifake-real-and-ai-generated-synthetic-images -p /content/data")
os.system("unzip -q /content/data/cifake-real-and-ai-generated-synthetic-images.zip -d /content/data/cifake")

# ============================================================
# Cell 3: Load and Explore Dataset
# ============================================================
train_transform = transforms.Compose([
    transforms.Resize(232),
    transforms.CenterCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

test_transform = transforms.Compose([
    transforms.Resize(232),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_dataset = datasets.ImageFolder("/content/data/cifake/train", transform=train_transform)
test_dataset = datasets.ImageFolder("/content/data/cifake/test", transform=test_transform)

print(f"Training samples: {len(train_dataset)}")
print(f"Test samples: {len(test_dataset)}")
print(f"Classes: {train_dataset.classes}")

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=2)

def show_samples(dataset, num_images=8):
    fig, axes = plt.subplots(2, num_images // 2, figsize=(14, 6))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    for i, ax in enumerate(axes.flat):
        img, label = dataset[i * (len(dataset) // num_images)]
        img = img.numpy().transpose(1, 2, 0)
        img = std * img + mean
        img = np.clip(img, 0, 1)
        ax.imshow(img)
        ax.set_title(dataset.classes[label])
        ax.axis("off")
    plt.tight_layout()
    plt.show()

show_samples(train_dataset)

# ============================================================
# Cell 4: Build the Model
# ============================================================
# Load pre-trained ResNet50
model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

# Freeze all backbone layers
for param in model.parameters():
    param.requires_grad = False

# Replace the final FC layer for binary classification
model.fc = nn.Linear(2048, 2)

# Move model to GPU
model = model.to(device)

# Define loss and optimizer (only train the new FC layer)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

print(f"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)")

Your model is assembled and ready to learn. In the next step, you will write the training loop that feeds batches of images through the model and teaches it to distinguish real photos from AI fakes.

Train the Model

Your model is built and ready to learn. Right now, the frozen ResNet50 backbone can extract features from images, but the new classification head has random weights. It has no idea what makes an image real or AI-generated.

In this step, you will write a training loop that teaches the model to distinguish between the two classes. The loop feeds batches of images through the network, measures how wrong the predictions are, and adjusts the weights to reduce that error.

In this step, get ready to:
  • Write the forward pass, loss calculation, and backpropagation training loop.
  • Add validation evaluation to track accuracy and detect overfitting.
  • Train for 5 epochs and confirm model accuracy improves each round.
Write the training loop

A training loop is the core cycle of deep learning. Each iteration (called an epoch) passes every training image through the model, computes how far off the predictions are, and nudges the weights in the right direction.

  • Create a new code cell in your notebook below Cell 4, then paste the following code:
# ============================================================
# Cell 5: Training Loop
# ============================================================
num_epochs = 5
train_losses = []
val_accuracies = []

for epoch in range(num_epochs):
    # Training phase
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(train_loader)
    train_losses.append(epoch_loss)

What does this code do?

This code sets up the training phase that repeats for 5 epochs.

  • model.train() switches the model to training mode, enabling features like dropout (not used here, but good practice).
  • optimizer.zero_grad() resets the gradients from the previous batch so they do not accumulate.
  • loss.backward() computes how much each weight contributed to the error (backpropagation).
  • optimizer.step() updates the weights using the computed gradients.

Why does this order matter?

The sequence zero_grad → forward → loss → backward → step is the standard PyTorch training recipe. Skipping any step (or reordering them) causes silent bugs like stale gradients or weights that never update.

You will not run this cell yet. The next substep adds validation evaluation to the same cell.

Add validation evaluation at the end of each epoch

Training loss tells you how well the model fits the training data, but it does not tell you if the model generalizes to new images. Validation evaluation measures accuracy on unseen data after every epoch so you can spot overfitting (the model memorizes training images instead of learning general patterns).

  • Still inside the same Cell 5, paste the following code directly below the last line you added. Make sure the indentation stays inside the outer for loop (one level of indent):
    # Validation phase
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    val_acc = 100 * correct / total
    val_accuracies.append(val_acc)
    print(f"Epoch [{epoch+1}/{num_epochs}] - Loss: {epoch_loss:.4f} - Val Accuracy: {val_acc:.2f}%")

What does this code do?

This validation block runs after each epoch's training finishes.

  • model.eval() switches the model to evaluation mode, disabling training-specific behavior.
  • torch.no_grad() disables gradient computation since you are only measuring, not learning. This saves GPU memory and speeds up inference.
  • torch.max(outputs, 1) picks the class (REAL or FAKE) with the highest prediction score for each image.
  • The final print statement gives you a progress report each epoch so you can watch accuracy climb.
Run training for 5 epochs
  • Click the play button on Cell 5 (or press Shift+Enter) to start training.

How long does training take?

With a T4 GPU, each epoch takes about 1-2 minutes. The full 5-epoch run finishes in under 10 minutes. If you are on CPU (no GPU available), expect 10-15 minutes per epoch.

You should see output printed after each epoch showing the loss decreasing and accuracy increasing, similar to this:

Epoch [1/5] - Loss: 0.1842 - Val Accuracy: 92.45%
Epoch [2/5] - Loss: 0.0934 - Val Accuracy: 94.80%
Epoch [3/5] - Loss: 0.0721 - Val Accuracy: 95.32%
Epoch [4/5] - Loss: 0.0598 - Val Accuracy: 95.78%
Epoch [5/5] - Loss: 0.0512 - Val Accuracy: 96.10%

What should I see?

Your exact numbers will differ slightly due to random shuffling, but the pattern should be the same. Training loss should decrease each epoch, and validation accuracy should rise above 95% by epoch 5.

If validation accuracy stops improving while training loss keeps dropping, that signals overfitting. With only 5 epochs and a frozen backbone, this is unlikely here.

Training seems stuck or accuracy is low?

  • Check that your runtime is connected to a GPU. Go to Runtime > Change runtime type and confirm T4 GPU is selected. Training on CPU is much slower and may appear stuck.
  • Verify the indentation of the validation phase code. It must be inside the outer for loop (one indent level) but outside the inner batch loop.
  • Make sure you ran all previous cells in order. The model, optimizer, and data loaders must already exist in memory.

My training loop is not producing expected results. Help me debug.

✔️ Awesome, I've got everything!

Great. Double check you see 5 lines of output with decreasing loss and increasing accuracy.

ⓧ I'd like to double check the full code

Here is the complete Cell 5 for comparison:

# ============================================================
# Cell 5: Training Loop
# ============================================================
num_epochs = 5
train_losses = []
val_accuracies = []

for epoch in range(num_epochs):
    # Training phase
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(train_loader)
    train_losses.append(epoch_loss)

    # Validation phase
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    val_acc = 100 * correct / total
    val_accuracies.append(val_acc)
    print(f"Epoch [{epoch+1}/{num_epochs}] - Loss: {epoch_loss:.4f} - Val Accuracy: {val_acc:.2f}%")

Your model is now trained and achieving 95%+ accuracy on unseen images. Next up, you will visualize those training curves and build an inference pipeline that classifies any new image you throw at it.

Evaluate and Visualize Results

Your model is trained and you've watched accuracy climb across 5 epochs. Numbers in a terminal are useful, but they don't tell the full story of how your model learned or where it still struggles.

In this step, you'll visualize the training process with loss curves, reveal exactly which images trip up the model using a confusion matrix, and build a reusable inference function that can classify any new image you throw at it.

In this step, get ready to:
  • Plot training loss and validation accuracy curves to see how the model learned over time.
  • Generate a confusion matrix showing exactly where the model gets predictions right and wrong.
  • Build an inference function that classifies any image as REAL or FAKE with a confidence score.
Plot training loss and validation accuracy curves

Training curves are the standard way ML engineers diagnose how a model learned. A loss curve that drops steadily means the model is learning. An accuracy curve that plateaus tells you when additional training stops helping.

  • Create a new code cell below your training loop cell.
  • Plot both the training loss and validation accuracy side by side by pasting this code:
# Plot training curves
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(range(1, num_epochs + 1), train_losses, marker="o")
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Training Loss")
ax1.set_title("Training Loss Over Epochs")

ax2.plot(range(1, num_epochs + 1), val_accuracies, marker="o", color="green")
ax2.set_xlabel("Epoch")
ax2.set_ylabel("Validation Accuracy (%)")
ax2.set_title("Validation Accuracy Over Epochs")
plt.tight_layout()
plt.show()

What does this code do?

This creates two side-by-side charts using matplotlib. The left chart plots train_losses across epochs so you can see loss decreasing. The right chart plots val_accuracies so you can confirm accuracy is climbing.

If the loss curve were going UP or the accuracy curve were going DOWN, that would signal a problem with training. Healthy curves show loss decreasing and accuracy increasing until they plateau.

  • Run the cell by pressing the play button or hitting Shift+Enter.
  • You should see two charts appear below the cell. The left chart shows training loss dropping from a higher value down toward zero. The right chart shows validation accuracy climbing toward 95% or higher.

Not seeing any plots?

Make sure you ran the training loop cell first (Cell 5). The train_losses and val_accuracies lists must be populated before you can plot them.

If you see an error about num_epochs not being defined, re-run the training cell. Variables from earlier cells are lost if your runtime disconnected.

Help me debug my training curve plot

Generate a confusion matrix on the test set

A confusion matrix shows you the full breakdown of your model's predictions. It reveals not just overall accuracy, but how many REAL images got misclassified as FAKE and vice versa.

  • Create a new code cell below the plotting cell.
  • Generate predictions on the entire test set and display a confusion matrix by pasting this code:
# Confusion matrix
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
    for images, labels in test_loader:
        images = images.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        all_preds.extend(predicted.cpu().numpy())
        all_labels.extend(labels.numpy())

cm = confusion_matrix(all_labels, all_preds)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=train_dataset.classes)
disp.plot(cmap="Blues")
plt.title("Confusion Matrix - AI Image Detector")
plt.show()

What does this code do?

First, the code runs every test image through the model and collects predictions into all_preds and ground-truth labels into all_labels. This happens inside torch.no_grad() since we don't need gradients for evaluation.

Then sklearn's confusion_matrix compares predictions against truth, and ConfusionMatrixDisplay renders it as a heatmap. The diagonal cells (top-left and bottom-right) show correct predictions. Off-diagonal cells show mistakes.

  • Run the cell by pressing Shift+Enter.
  • You should see a blue heatmap grid with FAKE and REAL labels on both axes. The diagonal squares should have large numbers (correctly classified images) and the off-diagonal squares should have small numbers (misclassifications).

Seeing low numbers on the diagonal?

If your diagonal numbers are close to the off-diagonal numbers, the model isn't classifying well. This usually means the training loop didn't complete properly. Go back to Cell 5 and re-run the training for all 5 epochs.

If you see a memory error, your runtime may have disconnected. Reconnect by clicking **Connect** in the top right, then re-run all cells from the top.

Help me understand my confusion matrix results

What is a confusion matrix telling me?

Each cell counts how many test images fell into that category. The top-left cell shows FAKE images correctly identified as FAKE. The bottom-right shows REAL images correctly identified as REAL.

The off-diagonal cells reveal the model's failure modes. A high number in the top-right means the model is calling FAKE images REAL. This kind of breakdown is more actionable than a single accuracy number.

Build an inference function for new images

Your model can classify images within the test set, but the real goal is classifying any new image you encounter. An inference function wraps the preprocessing and prediction steps into a single reusable call.

  • Create a new code cell below the confusion matrix cell.
  • Build the inference function by pasting this code:
from PIL import Image

def predict_image(image_path, model, transform, class_names):
    """Classify a single image as REAL or FAKE."""
    model.eval()
    image = Image.open(image_path).convert("RGB")
    input_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(input_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        confidence, predicted_class = torch.max(probabilities, 1)

    prediction = class_names[predicted_class.item()]
    conf = confidence.item() * 100
    print(f"Prediction: {prediction} (Confidence: {conf:.1f}%)")
    return prediction, conf

# Example usage:
# predict_image("/content/test_image.jpg", model, test_transform, train_dataset.classes)

What does this code do?

The predict_image function takes a file path to any image and runs it through the full classification pipeline. It opens the image with PIL, applies the same transforms the model was trained on, and feeds it to the model.

The torch.softmax call converts raw model outputs into probabilities (0 to 1). Then torch.max picks the class with the highest probability and reports it as the confidence score.

  • Run the cell by pressing Shift+Enter. The function is now defined and ready to use.
  • Test it on an image from the test set by creating one more cell and running this:
# Test on a sample image from the test set
sample_path = test_dataset.samples[0][0]
predict_image(sample_path, model, test_transform, train_dataset.classes)
  • You should see output like Prediction: FAKE (Confidence: 99.2%) or similar, with a class name and a high confidence percentage.

Seeing a file not found error?

Make sure you ran the dataset loading cell (Cell 3) earlier in your session. The test_dataset.samples list references files in /content/data/cifake/ which only exists after the download and extraction cell runs.

Help me fix my predict_image function

Double check that your evaluation and inference code matches the full notebook so far.

✔️ Awesome, I've got everything!

Great. Make sure all your cells have been run in order and your outputs are visible.

ⓧ I'd like to double check the full code

Here is the complete notebook code from Cell 1 through Cell 7. Compare it against your cells to make sure nothing is missing.

# AI-Generated Image Detector - Google Colab Notebook
# Fine-tunes ResNet50 on CIFAKE dataset to classify real vs AI-generated images

# ============================================================
# Cell 1: Setup and Imports
# ============================================================
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
from torchvision import datasets, transforms
from torchvision.models import resnet50, ResNet50_Weights
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

# Verify GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ============================================================
# Cell 2: Download CIFAKE Dataset
# ============================================================
os.makedirs("/root/.kaggle", exist_ok=True)
os.rename("/content/kaggle.json", "/root/.kaggle/kaggle.json")
os.chmod("/root/.kaggle/kaggle.json", 0o600)

os.system("kaggle datasets download -d birdy654/cifake-real-and-ai-generated-synthetic-images -p /content/data")
os.system("unzip -q /content/data/cifake-real-and-ai-generated-synthetic-images.zip -d /content/data/cifake")

# ============================================================
# Cell 3: Load and Explore Dataset
# ============================================================
train_transform = transforms.Compose([
    transforms.Resize(232),
    transforms.CenterCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

test_transform = transforms.Compose([
    transforms.Resize(232),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_dataset = datasets.ImageFolder("/content/data/cifake/train", transform=train_transform)
test_dataset = datasets.ImageFolder("/content/data/cifake/test", transform=test_transform)

print(f"Training samples: {len(train_dataset)}")
print(f"Test samples: {len(test_dataset)}")
print(f"Classes: {train_dataset.classes}")

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=2)

def show_samples(dataset, num_images=8):
    fig, axes = plt.subplots(2, num_images // 2, figsize=(14, 6))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    for i, ax in enumerate(axes.flat):
        img, label = dataset[i * (len(dataset) // num_images)]
        img = img.numpy().transpose(1, 2, 0)
        img = std * img + mean
        img = np.clip(img, 0, 1)
        ax.imshow(img)
        ax.set_title(dataset.classes[label])
        ax.axis("off")
    plt.tight_layout()
    plt.show()

show_samples(train_dataset)

# ============================================================
# Cell 4: Build the Model
# ============================================================
model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

for param in model.parameters():
    param.requires_grad = False

model.fc = nn.Linear(2048, 2)
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

print(f"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")

# ============================================================
# Cell 5: Training Loop
# ============================================================
num_epochs = 5
train_losses = []
val_accuracies = []

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_loss = running_loss / len(train_loader)
    train_losses.append(epoch_loss)

    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    val_acc = 100 * correct / total
    val_accuracies.append(val_acc)
    print(f"Epoch [{epoch+1}/{num_epochs}] - Loss: {epoch_loss:.4f} - Val Accuracy: {val_acc:.2f}%")

# ============================================================
# Cell 6: Evaluate and Visualize
# ============================================================
# Plot training curves
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(range(1, num_epochs + 1), train_losses, marker="o")
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Training Loss")
ax1.set_title("Training Loss Over Epochs")

ax2.plot(range(1, num_epochs + 1), val_accuracies, marker="o", color="green")
ax2.set_xlabel("Epoch")
ax2.set_ylabel("Validation Accuracy (%)")
ax2.set_title("Validation Accuracy Over Epochs")
plt.tight_layout()
plt.show()

# Confusion matrix
model.eval()
all_preds = []
all_labels = []
with torch.no_grad():
    for images, labels in test_loader:
        images = images.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        all_preds.extend(predicted.cpu().numpy())
        all_labels.extend(labels.numpy())

cm = confusion_matrix(all_labels, all_preds)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=train_dataset.classes)
disp.plot(cmap="Blues")
plt.title("Confusion Matrix - AI Image Detector")
plt.show()

# ============================================================
# Cell 7: Inference Function
# ============================================================
from PIL import Image

def predict_image(image_path, model, transform, class_names):
    """Classify a single image as REAL or FAKE."""
    model.eval()
    image = Image.open(image_path).convert("RGB")
    input_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(input_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        confidence, predicted_class = torch.max(probabilities, 1)

    prediction = class_names[predicted_class.item()]
    conf = confidence.item() * 100
    print(f"Prediction: {prediction} (Confidence: {conf:.1f}%)")
    return prediction, conf

# Example usage:
# predict_image("/content/test_image.jpg", model, test_transform, train_dataset.classes)

You now have a complete AI image detection pipeline. Your model can classify images, you can visualize how it learned, and your inference function makes it easy to test on any new image you find online.

Secret mission

Add Grad-CAM Visualization

Your model can classify images as real or AI-generated with high accuracy. But can you see what it's actually looking at? In this secret mission, you'll implement Grad-CAM (Gradient-weighted Class Activation Mapping) to generate heatmaps that reveal which regions of an image your model focuses on when making its decision. Research on the CIFAKE dataset shows that models often focus on small background imperfections rather than the main subject.

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

This project runs entirely on Google Colab's free tier, so there are no ongoing costs. The CIFAKE dataset and trained model live in Colab's temporary storage and are automatically deleted when your runtime disconnects.

Resources you used:

  • Google Colab notebook (saved in your Google Drive).
  • Google Colab GPU runtime (holds dataset and model in temporary storage).
  • Kaggle account (used to download the CIFAKE dataset).

Keep everything running

No action needed. Choose this if you're still actively experimenting with your model or want to try classifying more images.

  • Your notebook stays in Google Drive and you can reopen it anytime.
  • Your Colab runtime will eventually disconnect on its own after a period of inactivity (up to 12 hours max session).
  • Your Kaggle account remains active for future dataset downloads.

Pause - I'll come back to this later

Disconnect your runtime to free up GPU resources for others. Your notebook code stays saved in Google Drive.

  • In your Colab notebook, click Runtime in the top menu bar.
  • Click Disconnect and delete runtime.
  • This deletes the temporary dataset and model from Colab's storage, but your notebook code remains in Google Drive.
  • When you return, reconnect to a GPU runtime and re-run all cells to retrain the model.

Delete - I don't want to use this again

Remove your notebook and disconnect the runtime completely.

  • In your Colab notebook, click Runtime in the top menu bar, then click Disconnect and delete runtime. This removes the dataset and model from temporary storage.
  • Go to Google Drive at drive.google.com.
  • Navigate to the Colab Notebooks folder.
  • Right-click your AI image detector notebook and select Move to trash.
  • Your Kaggle account is free and can be kept for future projects. No action needed there.

Nice Work!

Nice Work!

You just trained a deep learning model to detect AI-generated images. From loading 120,000 images to running inference on new photos, you completed the full transfer learning workflow used in production ML systems.

Here's what you learned:

  • Applied transfer learning by fine-tuning a pre-trained ResNet50 model to classify images as real or AI-generated.
  • Wrote a complete training loop with validation monitoring that achieved 95%+ accuracy on the CIFAKE dataset.
  • Evaluated model performance using confusion matrices and built an inference pipeline to classify any new image with a confidence score.
  • Added Grad-CAM visualization to see which regions of an image the model focuses on when making its real vs. fake decision.

Ready to quiz yourself?