Skip to content

Train Machine Learning Models on Colab GPU

Google Colab

Google Colab enables you to run Jupyter notebooks in the cloud with the option to use a CPU or accelerate computations by adding GPU or TPU support. We will use the free Colab tier, but for longer training jobs or access to better GPUs (e.g., T4, P100, or V100), the paid Colab Pro or Colab Pro+ option may be a better choice. Navigate to Colab website and check out an example Jupyter notebook that uses a GPU for machine learning training.

Free Colab Limitations

In the free version of Colab, access to expensive resources like GPUs is heavily restricted and depends on availability as well as your usage patterns. Additionally, notebooks can run for a maximum of 12 hours in a single session, and the types of GPUs and TPUs available in Colab may vary over time. Paid plans offer increased compute availability based on your compute unit balance.

One advantage of using Colab is that common machine learning packages like keras, tensorflow, and xgboost come pre-installed with GPU support, so there is no need for Python environment setup as required on on-premise systems like Sherlock. Colab also integrates seamlessly with Google Drive, allowing you to upload your data and Jupyter notebooks to Drive for use in the Colab environment.

Start a Colab Notebook

Follow these steps to open your notebook in Colab:

  1. Upload your notebook to Google Drive.
  2. Right-click on the notebook file.
  3. Select Open with from the context menu.
  4. Choose Google Colaboratory.

If Google Colaboratory does not appear in the drop-down menu:

  1. Click Connect more apps at the bottom of the Open with menu.
  2. Search for Colaboratory in the app store.
  3. Select Colaboratory and connect it to your Google account.

Once connected, Colab will be available as an option in the Open with menu for future use. This only needs to be done once. After that, Colab will be available as an option for opening notebooks in the future.

If you need to create a notebook from scratch:

  1. Go to your Google Drive.
  2. Sign in with your Google account, if you’re not already signed in.
  3. Click + New > More > Google Colaboratory to start a new notebook.
  4. Save the notebook to Google Drive by clicking File > Save and selecting a location in your Drive.

Switch to Using a GPU

By default, notebooks in Colab run on the CPU. To leverage faster computations, you can switch to a GPU or TPU. Follow these steps to select a GPU accelerator:

  1. Go to Edit -> Notebook settings in the Colab interface.
  2. From the Hardware accelerator drop-down menu, select GPU (or TPU if your task is TPU-optimized).
  3. Click Save.

Now your notebook will run on the GPU. GPU-accelerated execution allows faster training and inference for deep learning models.

Access Data from Google Drive

You can access data stored in your Google Drive from Colab by mounting it. This allows you to easily upload, store, and load files directly from Google Drive. To do this, run the following code in a cell within a Colab notebook (e.g., colab_mnist.ipynb):

colab_mnist.ipynb
# Mount Google Drive
from google.colab import drive
drive.mount('/content/drive')
When prompted, click Connect to Google Drive, then choose your Google account to authorize access for your project. Once mounted, your Google Drive is accessible at /content/drive. For example, if you have a folder named ml_project in your Google Drive with a subdirectory data, you can reference it in Colab like this:

colab_mnist.ipynb
# Define input data path
path_to_data = '/content/drive/MyDrive/ml_project/data'

Run the Colab Notebook

Now we are ready to run the notebook, cell by cell. The following Python code will train a simple MNIST convolutional neural network (ConvNet):

colab_mnist.ipynb
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import numpy as np
from tensorflow import keras
from tensorflow.keras import layers

# Model / data parameters
num_classes = 10
input_shape = (28, 28, 1)

# the data, split between train and test sets
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# Scale images to the [0, 1] range
x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255

# Make sure images have shape (28, 28, 1)
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)
print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")

# convert class vectors to binary class matrices
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

# build the model
model = keras.Sequential(
   [
       keras.Input(shape=input_shape),
       layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
       layers.MaxPooling2D(pool_size=(2, 2)),
       layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
       layers.MaxPooling2D(pool_size=(2, 2)),
       layers.Flatten(),
       layers.Dropout(0.5),
       layers.Dense(num_classes, activation="softmax"),
   ]
)

print(model.summary())

# train the model
batch_size = 128
epochs = 15

model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])

model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1)

# Evaluate the trained model
score = model.evaluate(x_test, y_test, verbose=0)
print("Test loss:", score[0])
print("Test accuracy:", score[1])

If you execute this code in a cell, the output should look similar to this:

colab_mnist.ipynb
x_train shape: (60000, 28, 28, 1)
60000 train samples
10000 test samples
Model: "sequential_1"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓
┃ Layer (type)                         ┃ Output Shape                ┃         Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩
│ conv2d_2 (Conv2D)                    │ (None, 26, 26, 32)          │             320 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ max_pooling2d_2 (MaxPooling2D)       │ (None, 13, 13, 32)          │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ conv2d_3 (Conv2D)                    │ (None, 11, 11, 64)          │          18,496 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ max_pooling2d_3 (MaxPooling2D)       │ (None, 5, 5, 64)            │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ flatten_1 (Flatten)                  │ (None, 1600)                │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ dropout_1 (Dropout)                  │ (None, 1600)                │               0 │
├──────────────────────────────────────┼─────────────────────────────┼─────────────────┤
│ dense_1 (Dense)                      │ (None, 10)                  │          16,010 │
└──────────────────────────────────────┴─────────────────────────────┴─────────────────┘
Total params: 34,826 (136.04 KB)
Trainable params: 34,826 (136.04 KB)
Non-trainable params: 0 (0.00 B)
None
Epoch 1/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 48s 112ms/step - accuracy: 0.7578 - loss: 0.7626 - val_accuracy: 0.9785 - val_loss: 0.0816
Epoch 2/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 79s 107ms/step - accuracy: 0.9645 - loss: 0.1147 - val_accuracy: 0.9817 - val_loss: 0.0633
Epoch 3/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 48s 114ms/step - accuracy: 0.9725 - loss: 0.0895 - val_accuracy: 0.9870 - val_loss: 0.0486
Epoch 4/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 79s 106ms/step - accuracy: 0.9780 - loss: 0.0712 - val_accuracy: 0.9900 - val_loss: 0.0414
Epoch 5/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 82s 106ms/step - accuracy: 0.9802 - loss: 0.0627 - val_accuracy: 0.9897 - val_loss: 0.0388
Epoch 6/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 80s 102ms/step - accuracy: 0.9828 - loss: 0.0590 - val_accuracy: 0.9910 - val_loss: 0.0344
Epoch 7/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 44s 105ms/step - accuracy: 0.9838 - loss: 0.0512 - val_accuracy: 0.9907 - val_loss: 0.0322
Epoch 8/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 83s 107ms/step - accuracy: 0.9844 - loss: 0.0507 - val_accuracy: 0.9915 - val_loss: 0.0301
Epoch 9/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 47s 112ms/step - accuracy: 0.9871 - loss: 0.0419 - val_accuracy: 0.9917 - val_loss: 0.0307
Epoch 10/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 45s 106ms/step - accuracy: 0.9866 - loss: 0.0428 - val_accuracy: 0.9918 - val_loss: 0.0293
Epoch 11/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 44s 105ms/step - accuracy: 0.9863 - loss: 0.0432 - val_accuracy: 0.9923 - val_loss: 0.0298
Epoch 12/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 45s 107ms/step - accuracy: 0.9883 - loss: 0.0361 - val_accuracy: 0.9920 - val_loss: 0.0301
Epoch 13/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 82s 107ms/step - accuracy: 0.9884 - loss: 0.0345 - val_accuracy: 0.9928 - val_loss: 0.0282
Epoch 14/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 82s 106ms/step - accuracy: 0.9896 - loss: 0.0321 - val_accuracy: 0.9920 - val_loss: 0.0285
Epoch 15/15
422/422 ━━━━━━━━━━━━━━━━━━━━ 82s 107ms/step - accuracy: 0.9882 - loss: 0.0360 - val_accuracy: 0.9923 - val_loss: 0.0275
Test loss: 0.023274529725313187
Test accuracy: 0.9919000267982483

Ensure that the paths for loading data and saving outputs and are correctly set as needed. Below is an example of how to save results, the trained model, and predictions to Google Drive:

colab_mnist.ipynb
import os

# Define the output directory (make sure it exists)
output_dir = './outputs'
os.makedirs(output_dir, exist_ok=True)

# Save test results to a file in the output directory
results_file = os.path.join(output_dir, 'results.txt')
with open(results_file, 'w') as f:
   f.write(f"Test loss: {score[0]}\n")
   f.write(f"Test accuracy: {score[1]}\n")
print(f"Evaluation results saved to {results_file}")

# Save predictions as a NumPy array in the output directory
predictions_file = os.path.join(output_dir, 'predictions.npy')
predictions = model.predict(x_test)
np.save(predictions_file, predictions)
print(f"Predictions saved to {predictions_file}")

# Save the entire model in the native Keras format
model_file = os.path.join(output_dir, 'mnist_model.keras')
model.save(model_file)
print(f"Model saved to {model_file}")

Running the above code will generate files in your Google Drive, with output messages like this:

colab_mnist.ipynb
Evaluation results saved to ./outputs/results.txt
Predictions saved to ./outputs/predictions.npy
Model saved to ./outputs/mnist_model.keras
313/313 ━━━━━━━━━━━━━━━━━━━━ 4s 14ms/step

Save Your Work

Idle sessions may disconnect from Colab after 90 minutes of inactivity. Reconnect by refreshing the page and remounting resources (like Google Drive) if needed. When the session times out or disconnects, autosave will retain your most recent changes. However, it’s a good practice to manually save your notebook periodically to ensure no work is lost.