# Building an Artist Attribution Model with PyTorch and ResNet-50

> Source: <https://dev.to/ghazy001/building-an-artist-attribution-model-with-pytorch-and-resnet-50-2pde>
> Published: 2026-06-27 22:21:53+00:00

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.

The project classifies paintings by artist.

Given an input image, the model predicts:

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.

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:

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.

The project includes several useful features:

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

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.

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 processingTo train the model locally, run:

```
python3 train.py
```

The training script handles several steps:

`artist_model.pth`

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

That makes the project easier to run across different machines.

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.

This project helped me understand several important deep learning concepts.

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.

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/
```

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.

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

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

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

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.