Building an Artist Attribution Model with PyTorch and ResNet-50

A developer built an artist attribution system using PyTorch and a pretrained ResNet-50 model. The system classifies paintings by artist with top-3 predictions, leveraging transfer learning for efficient training. The project includes scripts for training and inference, and is designed to run locally or in cloud environments.

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.