Building an Artist Attribution Model with PyTorch and ResNet-50

Introduction

Can an AI model look at a painting and predict who created it?

That was the main idea behind this project: an artist attribution system built with PyTorch and ResNet-50. The goal was to train a deep learning model on a dataset of paintings and allow it to predict the most likely artist behind a new image.

Instead of building a convolutional neural network from scratch, I used transfer learning with a pretrained ResNet-50 model. This made the project more practical, faster to train, and easier to adapt to a custom art dataset.

In this article, I’ll walk through the idea, project structure, training process, inference flow, and the lessons learned while building it.

What the Project Does

The project classifies paintings by artist.

Given an input image, the model predicts:

• the most likely artist
• the confidence score
• the top 3 possible artist guesses

Example output:

🎨 Predicted artist: Vincent van Gogh
🔒 Confidence: 0.87
🔎 Top 3 guesses:
 - Vincent van Gogh (0.874)
 - Claude Monet (0.054)
 - Paul Cézanne (0.032)

This makes the project useful as a practical introduction to computer vision, fine-tuning, and image classification.

Why ResNet-50?

ResNet-50 is a powerful convolutional neural network architecture that has already learned useful visual patterns from ImageNet.

For this project, using a pretrained model made sense because painting classification requires the model to understand visual features such as:

• brush strokes
• color palettes
• composition
• texture
• artistic style
• recurring visual patterns

Training a deep model from zero would require a much larger dataset and more compute. Transfer learning allowed me to reuse the knowledge from ResNet-50 and adapt it to the artist classification task.

Main Features

The project includes several useful features:

• Transfer learning with ResNet-50
• Safe image loading to skip corrupted images
• Support for CUDA, Apple Silicon M-series, and CPU
• Training and inference scripts
• Top-3 prediction output
• Colab notebook support for cloud training

This makes the project flexible enough to run locally or in a cloud notebook environment.

Project Structure

The repository is organized like this:

artist-classification/
├── dataset/              # dataset folder
├── train.py              # training script
├── predict.py            # inference script
├── artist_model.pth      # trained model weights
├── README.md
└── test.jpg

The dataset is not included directly in the repository because of its size. After downloading it, the expected folder structure is:

dataset/
├── train/
│   ├── artist_1/
│   ├── artist_2/
│   └── ...
└── val/
    ├── artist_1/
    ├── artist_2/
    └── ...

This structure is important because image classification tools in PyTorch often rely on folder names as class labels.

For example, if the folder is named Vincent_van_Gogh, the model can treat that folder as one class.

Installing Requirements

The project uses a simple Python setup.

Install the required dependencies:

pip install torch torchvision pillow

The main libraries are:

• torch for deep learning
• torchvision for pretrained models and image transforms
• Pillow for image loading and processing

Training the Model

To train the model locally, run:

python3 train.py

The training script handles several steps:

1. Detects the available device
2. Loads the painting dataset
3. Applies image transformations
4. Fine-tunes ResNet-50
5. Saves the trained model as artist_model.pth

One detail I liked about this project is that it supports multiple environments:

• CUDA for NVIDIA GPUs
• MPS for Apple Silicon Macs
• CPU as a fallback

That makes the project easier to run across different machines.

Running Inference

After training, you can classify a test image with:

python3 predict.py

The prediction script loads the trained model and runs inference on an image.

Instead of only returning one answer, it gives the top 3 predictions. This is useful because art attribution can be uncertain. Two artists may have similar visual styles, especially if they belong to the same movement or period.

A Top-3 output gives more context than a single prediction.

What I Learned

This project helped me understand several important deep learning concepts.

1. Transfer learning saves time

Starting from a pretrained ResNet-50 model made the project much more realistic. The model already understands general image features, so the training process focuses on adapting those features to paintings.

2. Dataset structure matters

For image classification, clean folder organization is very important. If the dataset is not structured correctly, training can fail or produce incorrect labels.

A simple structure like this works well:

train/class_name/
val/class_name/

3. Inference should be user-friendly

A model prediction is more useful when it includes confidence scores and alternative guesses.

Instead of only printing:

Vincent van Gogh

the project prints:

Vincent van Gogh (0.874)
Claude Monet (0.054)
Paul Cézanne (0.032)

This makes the output easier to understand and debug.

4. Hardware support improves accessibility

Not everyone has the same machine. Supporting CUDA, Apple Silicon, and CPU makes the project easier for more developers to try.

Possible Improvements

There are several ways this project could be improved in the future:

• Add a Streamlit or Gradio web interface
• Include evaluation metrics such as accuracy and confusion matrix
• Add data augmentation for better generalization
• Support batch prediction for multiple images
• Add model checkpointing during training
• Deploy the model as an API
• Compare ResNet-50 with EfficientNet or Vision Transformers

A simple web interface would make the project especially impressive because users could upload a painting and instantly see the predicted artist.

Final Thoughts

This project was a great way to combine art and machine learning.

By using PyTorch, transfer learning, and ResNet-50, I built a model that can classify paintings by artist and return the top predictions with confidence scores.

The most important lesson is that deep learning projects do not always need to start from scratch. With transfer learning, we can build useful and interesting computer vision applications faster while still learning the core ideas behind model training, inference, and evaluation.

If you are learning PyTorch or computer vision, building an artist classification model is a fun and practical project to try.