Keras vs. PyTorch: Alien vs. Predator recognition with transfer learning

In our previous post, we gave you an overview of the differences between Keras and PyTorch, aiming to help you pick the framework that’s better suited to your needs. Now, it’s time for a trial by combat. We’re going to pit Keras and PyTorch against each other, showing their strengths and weaknesses in action. We present a real problem, a matter of life-and-death: distinguishing Aliens from Predators!

Predator and Alien
Image taken from our dataset. Both Predator and Alien are deeply interested in AI.

We perform image classification, one of the computer vision tasks deep learning shines at. As training from scratch is unfeasible in most cases (as it is very data hungry), we perform transfer learning using ResNet-50 pre-trained on ImageNet. We get as practical as possible, to show both the conceptual differences and conventions.

At the same time we keep the code fairly minimal, to make it clear and easy to read and reuse. See notebooks on GitHub, Kaggle kernels or Neptune versions with fancy charts.

Wait, what’s transfer learning? And why ResNet-50?

In practice, very few people train an entire Convolutional Network from scratch (with random initialization), because it is relatively rare to have a dataset of sufficient size. Instead, it is common to pretrain a ConvNet on a very large dataset (e.g. ImageNet, which contains 1.2 million images with 1000 categories), and then use the ConvNet either as an initialization or a fixed feature extractor for the task of interest.– Andrej Karpathy (Transfer Learning – CS231n Convolutional Neural Networks for Visual Recognition)

Transfer learning is a process of making tiny adjustments to a network trained on a given task to perform another, similar task. In our case we work with the ResNet-50 model trained to classify images from the ImageNet dataset. It is enough to learn a lot of textures and patterns that may be useful in other visual tasks, even as alien as this Alien vs. Predator case. That way, we use much less computing power to achieve much better result.

In our case we do it the simplest way:

  • keep the pre-trained convolutional layers (so-called feature extractor), with their weights frozen,
  • remove the original dense layers, and replace them with brand-new dense layers we will use for training.

So, which network should be chosen as the feature extractor?

ResNet-50 is a popular model for ImageNet image classification (AlexNet, VGG, GoogLeNet, Inception, Xception are other popular models). It is a 50-layer deep neural network architecture based on residual connections, which are connections that add modifications with each layer, rather than completely changing the signal.

ResNet was the state-of-the-art on ImageNet in 2015. Since then, newer architectures with higher scores on ImageNet have been invented. However, they are not necessarily better at generalizing to other datasets (see the Do Better ImageNet Models Transfer Better? arXiv paper).

Ok, it’s time to dive into the code.

Let the match begin!

We do our Alien vs. Predator task in seven steps:

  1. Prepare the dataset
  2. Import dependencies
  3. Create data generators
  4. Create the network
  5. Train the model
  6. Save and load the model
  7. Make predictions on sample test images

We supplement this blog post with Python code in Jupyter Notebooks (Keras-ResNet50.ipynb, PyTorch-ResNet50.ipynb). This environment is more convenient for prototyping than bare scripts, as we can execute it cell by cell and peak into the output.

All right, let’s go!

0. Prepare the dataset

We created a dataset by performing a Google Search with the words “alien” and “predator”. We saved JPG thumbnails (around 250×250 pixels) and manually filtered the results. Here are some examples:

We split our data into two parts:

  • Training data (347 samples per class) – used for training the network.
  • Validation data (100 samples per class) – not used during the training, but needed in order to check the performance of the model on previously unseen data.

Keras requires the datasets to be organized in folders in the following way:

If you want to see the process of organizing data into directories, check out the data_prep.ipynb file. You can download the dataset from Kaggle.

1. Import dependencies

First, the technicalities. We assume that you have Python 3.5+, Keras 2.2.2 (with TensorFlow 1.10.1 backend) and PyTorch 0.4.1. Check out the requirements.txt file in the repo.

So, first, we need to import the required modules. We separate the code in Keras, PyTorch and common (one required in both).




We can check the frameworks’ versions by typing keras.__version__  and torch.__version__, respectively.

2. Create data generators

Normally, the images can’t all be loaded at once, as doing so would be too much for the memory to handle. At the same time, we want to benefit from the GPU’s performance boost by processing a few images at once. So we load images in batches (e.g. 32 images at once) using data generators. Each pass through the whole dataset is called an epoch.

We also use data generators for preprocessing: we resize and normalize images to make them as ResNet-50 likes them (224 x 224 px, with scaled color channels). And last but not least, we use data generators to randomly perturb images on the fly:

Performing such changes is called data augmentation. We use it to show a neural network which kinds of transformations don’t matter. Or, to put it another way, we train on a potentially infinite dataset by generating new images based on the original dataset.

Almost all visual tasks benefit, to varying degrees, from data augmentation for training. For more info about data augmentation, see as applied to plankton photos or how to use it in Keras. In our case, we randomly shear, zoom and horizontally flip our aliens and predators.

Here we create generators that:

  • load data from folders,
  • normalize data (both train and validation),
  • augment data (train only).



In Keras, you get built-in augmentations and preprocess_input  method normalizing images put to ResNet-50, but you have no control over their order. In PyTorch, you have to normalize images manually, but you can arrange augmentations in any way you like.

There are also other nuances: for example, Keras by default fills the rest of the augmented image with the border pixels (as you can see in the picture above) whereas PyTorch leaves it black. Whenever one framework deals with your task much better than the other, take a closer look to see if they perform preprocessing identically; we bet they don’t.

3. Create the network

The next step is to import a pre-trained ResNet-50 model, which is a breeze in both cases. We freeze all the ResNet-50’s convolutional layers, and only train the last two fully connected (dense) layers. As our classification task has only 2 classes (compared to 1000 classes of ImageNet), we need to adjust the last layer.

Here we:

  • load pre-trained network, cut off its head and freeze its weights,
  • add custom dense layers (we pick 128 neurons for the hidden layer),
  • set the optimizer and loss function.



We load the ResNet-50 from both Keras and PyTorch without any effort. They also offer many other well-known pre-trained architectures: see Keras’ model zoo and PyTorch’s model zoo.  So, what are the differences?

In Keras we may import only the feature-extracting layers, without loading extraneous data ( include_top=False). We then create a model in a functional way, using the base model’s inputs and outputs. Then we use  model.compile(...) to bake into it the loss function, optimizer and other metrics.

In PyTorch, the model is a Python object. In the case of models.resnet50, dense layers are stored in model.fc attribute. We overwrite them. The loss function and optimizers are separate objects. For the optimizer, we need to explicitly pass a list of parameters we want it to update.

Predator's wrist computer
Frame from ‘AVP: Alien vs. Predator’: Predators’ wrist computer. We’re pretty sure Predator could use it to compute logsoftmax.

In PyTorch, we should explicitly specify what we want to load to the GPU using .to(device) method. We have to write it each time we intend to put an object on the GPU, if available. Well…

Layer freezing works in a similar way. However, in The Batch Normalization layer of Keras is broken (as of the current version; thx Przemysław Pobrotyn for bringing this issue). That is – some layers get modified anyway, even with  trainable = False.

Keras and PyTorch deal with log-loss in a different way.

In Keras, a network predicts probabilities (has a built-in softmax function), and its built-in cost functions assume they work with probabilities.

In PyTorch we have more freedom, but the preferred way is to return logits. This is done for numerical reasons, performing softmax then log-loss means doing unnecessary log(exp(x)) operations. So, instead of using softmax, we use LogSoftmax (and NLLLoss) or combine them into one nn.CrossEntropyLoss  loss function.

4. Train the model

OK, ResNet is loaded, so let’s get ready to space rumble!

Predators' mother ship
Frame from ‘AVP: Alien vs. Predator’: the Predators’ Mother Ship. Yes, we’ve heard that there are no rumbles in space, but nothing is impossible for Aliens and Predators.

Now, we proceed to the most important step – model training. We need to pass data, calculate the loss function and modify network weights accordingly. While we already had some differences between Keras and PyTorch in data augmentation, the length of code was similar. For training… the difference is massive. Let’s see how it works!

Here we:

  • train the model,
  • measure the loss function (log-loss) and accuracy for both training and validation sets.



In Keras, the model.fit_generator performs the training… and that’s it! Training in Keras is just that convenient. And as you can find in the notebook, Keras also gives us a progress bar and a timing function for free. But if you want to do anything nonstandard, then the pain begins…

Predators' shuriken
Predator’s shuriken returning to its owner automatically. Would you prefer to implement its tracking ability in Keras or PyTorch?

PyTorch is on the other pole. Everything is explicit here. You need more lines to construct the basic training, but you can freely change and customize all you want.

Let’s shift gears and dissect the PyTorch training code. We have nested loops, iterating over:

  • epochs,
  • training and validation phases,
  • batches.

The epoch loop does nothing but repeat the code inside. The training and validation phases are done for three reasons:

  • Some special layers, like batch normalization (present in ResNet-50) and dropout (absent in ResNet-50), work differently during training and validation. We set their behavior by model.train() and model.eval(), respectively.
  • We use different images for training and for validation, of course.
  • The most important and least surprising thing: we train the network during training only. The magic commands optimizer.zero_grad(), loss.backward() and  optimizer.step() (in this order) do the job. If you know what backpropagation is, you appreciate their elegance.

We take care of computing the epoch losses and prints ourselves.

5. Save and load the model


Once our network is trained, often with high computational and time costs, it’s good to keep it for later. Broadly, there are two types of savings:

  • saving the whole model architecture and trained weights (and the optimizer state) to a file,
  • saving the trained weights to a file (keeping the model architecture in the code).
    It’s up to you which way you choose.

Here we:

  • save the model.



Alien is evolving
Frame from ‘Alien: Resurrection’: Alien is evolving, just like PyTorch.

One line of code is enough in both frameworks. In Keras you can either save everything to a HDF5 file or save the weights to HDF5 and the architecture to a readable json file. By the way: you can then load the model and run it in the browser.

Currently, PyTorch creators recommend saving the weights only. They discourage saving the whole model because the API is still evolving.


Loading models is as simple as saving. You should just remember which saving method you chose and the file paths.

Here we:

  • load the model.



In Keras we can load a model from a JSON file, instead of creating it in Python (at least when we don’t use custom layers). This kind of serialization makes it convenient for transfering models.

PyTorch can use any Python code. So pretty much we have to re-create a model in Python.

Loading model weights is similar in both frameworks.

6. Make predictions on sample test images

All right, it’s finally time to make some predictions! To fairly check the quality of our solution, we ask the model to predict the type of monsters from images not used for training. We can use the validation set, or any other image.

Here we:

  • load and preprocess test images,
  • predict image categories,
  • show images and predictions.





Prediction, like training, works in batches (here we use a batch of 3; though we could surely also use a batch of 1). In both Keras and PyTorch we need to load and preprocess the data. A rookie mistake is to forget about the preprocessing step (including color scaling). It is likely to work, but result in worse predictions (since it effectively sees the same shapes but with different colors and contrasts).

In PyTorch there are two more steps, as we need to:

  • convert logits to probabilities,
  • transfer data to the CPU and convert to NumPy (fortunately, the error messages are fairly clear when we forget this step).

And this is what we get:

It works!

And how about other images? If you can’t come up with anything (or anyone) else, try using photos of your co-workers. 🙂


As you can see, Keras and PyTorch differ significantly in terms of how standard deep learning models are defined, modified, trained, evaluated, and exported. For some parts it’s purely about different API conventions, while for others fundamental differences between levels of abstraction are involved.

Keras operates on a much higher level of abstraction. It is much more plug&play, and typically more succinct, but at the cost of flexibility.

PyTorch provides more explicit and detailed code. In most cases it means debuggable and flexible code, with only small overhead. Yet, training is way-more verbose in PyTorch. It hurts, but at times provides a lot of flexibility.

Transfer learning is a big topic. Try tweaking your parameters (e.g. dense layers, optimizer, learning rate, augmentation) or choose a different network architecture.

Have you tried transfer learning for image recognition? Consider the list below for some inspiration:

Pick Keras or PyTorch, choose a dataset and let us know how it went in the comments section below 🙂

Related Posts

10 replies
  1. Jason Do
    Jason Do says:

    Hello, I’m testing this out on my machine but I’m getting a weird error and I was wondering if you could help me fix it. When training the model with PyTorch, I’m getting the error: “Assertion `cur_target >= 0 && cur_target < n_classes' failed." From what I've read, this is due to some mismatch between 0 and 1 indexing. I've only copied and pasted your code and haven't changed anything but I still get this error. However, when I change the line nn.Linear(128, 2)).to(device) to nn.Linear(128, 3)).to(device), the training proceeds. That doesn't particularly make sense as I'm assuming the 2 refers to 2 classes. Do you have any idea what is wrong? I'm using PyTorch 0.4.1 as well.

    • Patryk Miziuła
      Patryk Miziuła says:

      Hi Jason,
      Did you clone the code from our repo? I’ve just tested the code and it works for me. What version of torchvision do you use? We recommend 0.2.1.

      • Jason Do
        Jason Do says:

        Hi Patryk,
        I didn’t clone the repo, but I copied and pasted the code snippets from both this website and your repo just to double check. All pieces of code appear to be the same/similar. I’m on python 3.6.3, pytorch 0.4.1, and torchvision 0.2.1, running on Jupyter Notebook with Anaconda Navigator. I still get the error, but I’ve worked around it by reducing to 3 as I said above and just changing the indices I’m pulling the percentages from. It works but I don’t understand why I’m getting the error in the first place.

        • Patryk Miziuła
          Patryk Miziuła says:

          Do you store the data exactly as we recommend (see point 0.)? datasets.ImageFolder needs to see exactly 2 subfolders to label the classes correctly.

          • Jason Do
            Jason Do says:

            My data is stored in data/train/alien, data/train/predator, data/validation/alien, and data/validation/predator as shown in the example. Like I said, the whole program works when I change the linear output to 3 and adjust accordingly, but I run into the out of bound error when the output is set to 2.

          • Patryk Miziuła
            Patryk Miziuła says:

            It looks like the categories were numbered as 1 and 2, not 0 and 1, pretty strange. Please run the following:
            inputs, labels = next(iter(dataloaders["train"]))
            outputs = F.softmax(model(inputs), dim=1)
            print(outputs.size(), labels.size())

            The expected output is: tensor [batch_size, 2] with rows summing up to 1 and tensor [batch_size] containing zeros and ones. If this is what you exactly get, but the problem still exists, then unfortunately I have no idea how to help you.

          • Jason Do
            Jason Do says:

            I’m replying here because I think we’ve reached the comment depth limit for the website. I ran your code snippet after I defined the optimizer and the criterion and before I trained the model. I got this as the output: Indeed, it looks like the categories are labeled as 1 and 2, though I do not see anywhere in your code where it does so. Do you have any idea of a fix?

          • Patryk Miziuła
            Patryk Miziuła says:

            One more idea before accusing PyTorch of being bugged: maybe you’ve got one more hidden folder in data/train? I added one manually on my computer (called it .sth) and then obtained labels being ones and twos, as in your case. By the way, does the Keras version work for you? Keras’ data generator also infers the number of classes from the folder content.

          • Jason Do
            Jason Do says:

            Holy cow. Well that solves that problem. There was a cheeky hidden auto-generated folder .ipynbcheckpoints chilling in my data folders, messing everything up. Thank you so much for being so responsive to such an ML newbie like me. You’ve been a tremendous help! Do you have any advice or resources for me to get more familiar with using pytorch? I’m trying to use your project as a base to classify chihuahuas and muffins. Also, which parts of your code can I tweak and fiddle with to fit different problems better?

          • Patryk Miziuła
            Patryk Miziuła says:

            I’m glad we finally nailed it, this was a good lesson for both of us.

            If you want to master PyTorch, I suggest you to simply dive into official PyTorch tutorials, they are great. We also plan to write the third (more advanced) blog post about Keras and PyTorch, but who knows when we’ll do that. Moreover, we sometimes conduct free webinars on ML/DL, so stay tuned. Finally, if you have some money to spend, consider our paid workshops.


Leave a Reply

Want to join the discussion?
Feel free to contribute!

Leave a Reply

Your email address will not be published. Required fields are marked *