Momo vs. Raccoon CNN

A pretrained ResNet-18 told my cat Momo apart from raccoons on every photo of a 50 photo test set, where a CNN trained from scratch scored 82%.

Code on GitHub

Close-up of Momo, a grey and cream striped cat, resting her chin on a cushion. — Momo, the cat the classifier has to tell apart from a raccoon.

Telling Momo from a raccoon

My family has a cat that looks suspiciously like a raccoon, complete with a striped tail and greyish fur. The easiest way to tell them apart, of course, is a convolutional neural network that picks up the differences in features between them and returns the probability that a photo shows our cat or a raccoon.

I built the project in PyTorch as an intro to deep learning. The data is 137 photos of Momo from our own phones and about 490 raccoon photos from Roboflow and iNaturalist. I split by photo group so near duplicates never land on both sides of the split, then ran every model on that same split. The test set is only 50 photos, so one photo is worth 2 points and the numbers below are rough.

A CNN from scratch

The first model has no preexisting weights. It stacks 8 convolution layers, 4 pooling layers, and a final average pooling layer, then a dropout layer and one linear layer, for 1.17 million parameters. It reached 78 to 82% accuracy at telling Momo from a raccoon, depending on the run.

Grad-CAM shows what a model looked at when it decided. On simple indoor photos the heatmaps sit on the cat. On busier photos they drift to the floor, the furniture, and the corners of the frame, and the confidence drops with them.

Diagram of the scratch CNN as a row of blocks, shrinking from a 128 by 128 input to an 8 by 8 feature map that gets thicker with each conv and pool stage, then ending in three thin layers. — The scratch CNN. Four blocks of two convolutions and a max pool take 3 × 128 × 128 down to 256 × 8 × 8, then average pooling, dropout, and one linear layer give a single raccoon logit.

Photos of Momo and raccoons beside scratch CNN Grad-CAM heatmaps, some on the animal and some on the background, with several raccoons labeled Momo. — Grad-CAM for the scratch CNN. Left photos are Momo, right photos are raccoons. All three raccoons shown here are misread as Momo.

Finetuning a pretrained ResNet

To improve on that, I tried a second approach: keep a ResNet-18 that was already trained on ImageNet, freeze most of it, and finetune the last block and the final layer on my photos. That trains 8.4 million of its 11.2 million parameters in a few minutes. It got every photo right on the 50 photo test set in both runs I did, and the Grad-CAM maps are much more reasonable.

Some false features still count toward the answer. A patch of floor or a sofa arm still lights up on some photos. More data, more unfrozen layers, or photos with more varied backgrounds should help with that.

Photos of Momo and raccoons beside ResNet-18 Grad-CAM heatmaps, with the red areas centered on the animals. — Grad-CAM for ResNet-18 on the same kind of photos. The strongest areas sit on the animal, with some bleed onto the background.

Gardening with a shotgun: the transformer approach

2017 was a big year for machine learning. The paper Attention Is All You Need introduced transformers, which now sit under every large language model. They train quickly because every position in the input can be processed in parallel, and their attention heads let each position weigh every other position when deciding what matters.

A vision transformer applies the same idea to images. It cuts the photo into 16 pixel patches, treats each patch like a word, and reads the answer from a class token that has attended to all of them. I used torchvision's ViT-B/16, pretrained on ImageNet, with 85.8 million parameters, of which I trained 14.2 million: the last two blocks, the final normalization layer, and a new output layer. For a cat and a raccoon this is gardening with a shotgun.

It also worked. The ViT got every photo right on the test set in both runs, and on 78 raccoon photos it had never seen it missed one. The ResNet missed 1 to 3 of them and the scratch CNN missed 7.

Why the ViT heatmaps look odd

The ViT heatmaps look strange next to the other two. Grad-CAM was designed around convolution feature maps, and a ViT only gives it a 14 by 14 grid of patches, so the maps come out blotchy and often light up the background. A different method such as attention rollout would suit it better, and I have not tried one yet.

Photos of Momo and raccoons beside ViT-B/16 Grad-CAM heatmaps made of blotchy patches, many of them on the background. — Grad-CAM for the ViT on the same photos. The maps are made of 16 pixel patches and often sit on the background even when the answer is right.

Lessons learned

Even in the era of the transformer, a convolutional network is still a data efficient way to predict in a very constrained setting. With fewer than 200 photos of Momo, the small CNN built from scratch got most of the way there.

Building on something pretrained did better, and in a more general way. Both pretrained models handled new raccoon photos better than the CNN I trained myself.

A high accuracy number can hide how the model got there. Grad-CAM showed all three models leaning on backgrounds at times, and with a 50 photo test set a single mistake moves the score by 2 points. Looking at the heatmaps told me more than the percentages did.

The raccoon photos come from Roboflow datasets and from iNaturalist observations under Creative Commons licenses. Claude helped me run the experiments and the comparisons.