Published in

Red Buffer

5 min readAug 20, 2021

Setting up PyTorch and TensorFlow on a Windows Machine

The most common operating system used with Deep Learning frameworks (PyTorch and TensorFlow)is Linux mainly because of extensive online support. However, Linux generally has a higher learning curve compared to windows making it difficult for novices to get started with Deep Learning.

The focus of this article is to explain the setup process of Deep Learning frameworks on a windows machine and make it simple for end-users.

The setup process consists of 4 main steps:

Installation of Nvidia Driver
Installation of Anaconda
Installation of CUDA Driver
Installation of Deep Learning Framework

Each step is explained in the article below.

** GeForce RTX 3060 (12 GB) was used for the creation of this guide.
** Python 3.6 was used for the creation of this guide.

Installation of Nvidia Driver

The first step in the Installation of Nvidia Driver is identifying the GPU installed on your machine. The machine that is in use has a GeForce RTX 3060 (12 GB) variant therefore we will be installing tools accordingly.

Now that we know the GPU model we can navigate to Nvidia Drivers Download Page.

On the download page add the relevant information of GPU that you are using. Once you have added the correct information press the search button as seen below:

Next, click the download button to start downloading the driver

Once the driver is downloaded, run the driver installer (.exe), and after installation is complete restart the machine for changes to occur.

Open Command Prompt (CMD) and type nvidia-smi to verify the installation of Nvidia Drivers.

The above output of nvidia-smi tells us the Nvidia Driver Version and the highest compatible CUDA version that is available.

For our case the versions are:

Driver Version: 471.21
CUDA Version: 11.4

Installation of Anaconda

We can install PyTorch and TensorFlow only using pip but that makes the task challenging. This article focuses on using Anaconda as conda makes the installation of packages easier by solving dependencies on its own without bothering the user.

The installation process is straight forward first navigate to the Anaconda website and download the Individual Edition. Once the download is complete run the setup executable file.

Create Conda Virtual Environment

It is recommended to create a virtual environment for each separate project so that there are no conflicts among the installed packages.

To create Conda Virtual Environment run the following command(s) in sequence on CMD:

conda activate
conda create -n [Environment Name] python=3.6
conda activate [Environment Name]

Installation of CUDA Driver

For this article, we will be installing CUDA using Anaconda as the conda package contains only the modules that are required to run PyTorch and TensorFlow code on GPU.

We will be installing CUDA 11.0 for this article however my Nvidia driver support is valid up to CUDA 11.4.

To install CUDA using conda run the following command(s) in CMD:

conda install cudatoolkit=11.0 cudnn

** Please install the appropriate CUDA version that your Nvidia Driver supports otherwise the package might be installed but would crash on execution

Installation of Deep Learning Framework

TensorFlow Installation

TensorFlow can either be installed using pip or anaconda. For now, we will be using pip for the installation of TensorFlow.

Before the installation process can be started we need to check the TensorFlow version that is compatible with the installed CUDA 11.0. The compatibility of TensorFlow with CUDA can be seen in the Compatibility Chart added on the Official TensorFlow Website.

By looking at the Compatibility Chart we see that with CUDA 11.0 we can install TensorFlow-GPU 2.4.0. For installation of TensorFlow-GPU 2.4.0 run the following command(s) in CMD:

pip install tensorflow-gpu==2.4

Once the installation is complete verify if the GPU is available for compute in your TensorFlow library run the following code snippet in IPython Console:

from tensorflow.python.client import device_lib

def get_available_gpus():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU']print(get_available_gpus())

The output of the above code snippet can be shown in the figure below and it is going to confirm that the installation of TensorFlow with GPU is successful.

** Fix indentation or formatting errors while copying the above snippet

PyTorch Installation

PyTorch can be installed using pip and Anaconda. For now, we will be using conda for the installation of PyTorch.

Before the installation process can be started we need to check the PyTorch version that is compatible with the installed CUDA 11.0. The compatibility of PyTorch with CUDA can be seen in the PyTorch Versions Page added on the Official PyTorch Website.

By looking at the Compatibility Chart we see that with CUDA 11.0 we can install PyTorch 1.7.0. For installation of PyTorch 1.7.0 run the following command(s) in CMD:

conda install pytorch==1.7.0 torchvision==0.8.0 -c pytorch

Once the installation is complete verify if the GPU is available for compute in your PyTorch library run the following code snippet in IPython Console:

import torch

torch.cuda.is_available()
>>> True

torch.cuda.current_device()
>>> 0

torch.cuda.device_count()
>>> 1

torch.cuda.get_device_name(0)
>>> 'NVIDIA GeForce RTX 3060'

If you are getting the same output on the IPython console as in the snippet above then the installation of PyTorch is successful.