Setting Up TensorFlow on a VPS For ML

By Anurag Singh

Updated on Oct 16, 2024

Setting Up TensorFlow on a VPS For ML

In this tutorial, we'll explain how to setting up TensorFlow on a VPS For ML.

TensorFlow is a powerful open-source machine learning framework used for various applications such as deep learning and neural networks. Setting it up on a Virtual Private Server (VPS) allows you to leverage the computational power for running deep learning models remotely. In this tutorial, we’ll walk through the step-by-step process of installing TensorFlow on a VPS.

Prerequisites:

  • A KVM VPS with Ubuntu 20.04 or higher (other Linux distributions should also work with minor adjustments)
  • A non-root user with sudo privileges
  • Basic knowledge of Python and machine learning
  • Adequate memory and storage for TensorFlow's resource demands

Step 1: Update the Server

Before we begin the installation, it's important to update the server's packages to the latest versions.

sudo apt update && sudo apt upgrade -y

This will ensure your VPS has all the latest security patches and software.

Step 2: Install Python and Required Dependencies

TensorFlow is written in Python, so we need to ensure that Python and the required build tools are installed.

Install Python and pip: TensorFlow requires Python 3.8 or higher, so we will install the necessary Python version.

sudo apt install python3-dev python3-pip python3-venv -y

Verify the Python installation: Check that Python and pip (Python package manager) have been installed correctly.

python3 --version
pip3 --version

Step 3: Set Up a Virtual Environment (Optional but Recommended)

Using a virtual environment allows you to isolate the TensorFlow installation, avoiding conflicts with other Python packages on your system.

Create a virtual environment: Navigate to your project directory and create a virtual environment named tensorflow_env.

python3 -m venv tensorflow_env

Activate the virtual environment:

source tensorflow_env/bin/activate

Once activated, you will see (tensorflow_env) in your shell, indicating that the environment is active.

Step 4: Install TensorFlow

With the virtual environment active, you can now install TensorFlow using pip.

Install TensorFlow: We will install the latest version of TensorFlow via pip. You can opt for either the CPU or GPU version, depending on whether your VPS supports GPU acceleration.

For CPU-only installation:

pip install tensorflow

For GPU installation: If your VPS has GPU support (like CUDA), you can install the GPU-optimized version.

pip install tensorflow-gpu

Verify TensorFlow installation: After installation, ensure TensorFlow is installed properly by checking its version.

python -c "import tensorflow as tf; print(tf.__version__)"

If the version number is displayed without errors, TensorFlow is successfully installed.

Step 5: Install Additional Libraries for Deep Learning

Deep learning models often require additional libraries for handling data, visualizations, and model evaluation. The most common ones are NumPy, Pandas, Matplotlib, and scikit-learn.

pip install numpy pandas matplotlib scikit-learn

Step 6: Example TensorFlow Model

Let’s create a simple TensorFlow script to verify that everything is working correctly. We’ll create a basic neural network model for classifying the MNIST dataset (a set of handwritten digits).

Create the TensorFlow script: Save the following script as

nano mnist_example.py:

Add following code:

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.datasets import mnist

# Load the MNIST dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# Normalize the data
x_train = x_train / 255.0
x_test = x_test / 255.0

# Define the input layer using Input()
input_layer = tf.keras.Input(shape=(28, 28))

# Build a simple neural network
model = models.Sequential([
    input_layer,
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(10, activation='softmax')
])

# Compile the model
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Train the model
model.fit(x_train, y_train, epochs=5)

# Evaluate the model
test_loss, test_acc = model.evaluate(x_test, y_test)

print(f'Test accuracy: {test_acc}')

Run the TensorFlow script: With TensorFlow installed, you can run the script to train the model on your VPS:

python mnist_example.py

The script will download the MNIST dataset, build a simple neural network, and train it. After training, it will print the test accuracy.

Step 7: Enabling GPU Acceleration (Optional)

If your VPS supports GPU acceleration, you can enable it for TensorFlow, which will significantly speed up model training, especially for deep learning models.

Install NVIDIA drivers: First, install the necessary NVIDIA drivers for CUDA.

sudo apt install nvidia-driver-470

Install CUDA toolkit: Install the CUDA toolkit and libraries.

sudo apt install nvidia-cuda-toolkit

Install cuDNN: Install cuDNN (CUDA Deep Neural Network library) for better GPU performance.

sudo apt install libcudnn8

Verify GPU support in TensorFlow: After installing the necessary drivers, verify that TensorFlow detects the GPU.

python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"

If the GPU is detected, the output will list available GPUs.

Step 8: Optimize Server for TensorFlow Workloads

Running deep learning models on a VPS can be resource-intensive. Here are some tips to optimize server performance for TensorFlow:

Increase Swap Memory: For memory-intensive operations, especially on CPU-bound servers, increasing swap memory can prevent out-of-memory (OOM) errors.

sudo fallocate -l 8G /swapfile
sudo chmod 600 /swapfile
sudo mkswap /swapfile
sudo swapon /swapfile

Use Batch Processing: While training deep learning models, use smaller batches to fit the data into memory efficiently.

Monitor Resource Usage: Use tools like htop to monitor CPU, memory, and GPU usage during training.

sudo apt install htop
htop

Step 9: Securing the TensorFlow Environment

Firewall Configuration: Ensure that your VPS is secure by configuring a firewall like UFW (Uncomplicated Firewall). Only allow essential ports.

sudo ufw allow OpenSSH
sudo ufw enable

Set Up SSH Access: Use SSH for remote access to your VPS, and disable root login for enhanced security.

Conclusion:

In this tutorial, we walked through the process of setting up TensorFlow on a VPS for running deep learning models. After updating the system, installing Python, setting up a virtual environment, and installing TensorFlow, we tested the installation with an example script. With this setup, you're ready to run machine learning models remotely on your dedicated server. Additionally, you can optimize for GPU acceleration if supported and secure the environment for production use.