Training Your Models on Cloud TPUs in 4 Easy Steps on Google Colab

Training Your Models on Cloud TPUs in 4 Easy Steps on Google Colab

You have a plain old TensorFlow model that’s too computationally expensive to train on your standard-issue work laptop. I get it. I’ve been there too, and if I’m being honest, seeing my laptop crash twice in a row after trying to train a model on it is painful to watch.

1 Like

Hey, extremely helpful article. I’ve been trying to use the tpu on google colab for some time. There are some bugs in your code.

import numpy as np
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical

(x_train, y_train), (x_test, y_test) = mnist.load_data()

# preprocessing the dataset
x_train = x_train.reshape(x_train.shape[0], 1) # 784, 1]
x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train /= 255
x_test /= 255

y_train = to_categorical(y_train, 10)
y_test = to_categorical(y_test, 10)
print (x_train.shape, y_train.shape)
print (x_test.shape, y_test.shape)

This returns errors, should be -

import numpy as np

import tensorflow as tf

from tensorflow.keras.datasets import mnist

from tensorflow.keras.utils import to_categorical

(x_train, y_train), (x_test, y_test) = mnist.load_data()

preprocessing the dataset

x_train = x_train.reshape(x_train.shape[0], 784) #784, 1]

x_train = x_train.astype(‘float32’)

x_test = x_test.reshape(x_test.shape[0], 784) #784, 1]

x_test = x_test.astype(‘float32’)

x_train /= 255

x_test /= 255

y_train = to_categorical(y_train, 10)

y_test = to_categorical(y_test, 10)

print (x_train.shape, y_train.shape)

print (x_test.shape, y_test.shape)

Hey there! Thank you for the kind words.

Looks like I accidentally made a typo. I’ll fix it asap. The thing about TensorFlow code is that things are deprecated or become old-fashioned rather quickly so articles may seem out of date at times :sweat_smile:

Thanks for pointing it out!

Best,
Rishabh

Ok that’s odd. It isn’t letting me edit. Here’s the fix, though:

import numpy as np
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical

(x_train, y_train), (x_test, y_test) = mnist.load_data()

# preprocessing the dataset
x_train = x_train.reshape([x_train.shape[0], 784])
x_train = x_train.astype(‘float32’)
x_test = x_test.reshape([x_test.shape[0], 784])
x_test = x_test.astype(‘float32’)
x_train /= 255
x_test /= 255

y_train = to_categorical(y_train, 10)
y_test = to_categorical(y_test, 10)

print (x_train.shape, y_train.shape)
print (x_test.shape, y_test.shape)
with strategy.scope():
	"""
	Note: This model can be whatever you want it to be.
	Here, I'm building a simple fully-connected network using 
	our distributed training strategy. 
	This essentailly takes our model and makes it 
	compatible to train on a TPU.
	"""
	model = tf.keras.models.Sequential()
	model.add(tf.keras.layers.Dense(512, input_shape=[784,], activation='relu'))
	model.add(tf.keras.layers.Dense(256, activation='relu'))
	model.add(tf.keras.layers.Dense(128, activation='relu'))
	model.add(tf.keras.layers.Dense(64, activation='relu'))
	model.add(tf.keras.layers.Dense(10, activation='softmax'))

	# compiling the model using the RMSProp optimizer 
	# and Sparse Categorical Crossentropy loss
	model.compile(
		optimizer=tf.train.RMSPropOptimizer(learning_rate=1e-2),
		loss=tf.keras.losses.sparse_categorical_crossentropy,
		metrics=['sparse_categorical_accuracy']
	)

	model.summary()

just want to fix for tensorflow default version in colab right now input_size change to input_shape

history = model.fit(x_train,  np.argmax(y_train, axis=1), epochs=20, steps_per_epoch=50)
model.save_weights('./mnist_model.h5', overwrite=True)

and also for debug this error now Can not squeeze dim[1], expected a dimension of 1, got 2
y_train convert to np.argmax(y_train, axis=1)