该环境是在Colab tensorflow 1.14.1 rc版本下的，由于版本比较新，但是API并不成熟和易用。如果旧有的代码无法正常使用的话，建议把tf重新安装为1.13.1版本(! pip install tensorflow==1.13.1)。

1. Fashion MNIST with Keras and TPUs

1.1 读取数据

使用tf.keras.datasets读取Fashion MNIST数据，其中需要把数据转换成float32类型：

import os

import tensorflow as tf
import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedShuffleSplit

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

x_train = x_train.astype('float32')
y_train = y_train.astype('float32')

x_test = x_test.astype('float32')
y_test = y_test.astype('float32')

# add empty color dimension
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

1.2 创建验证集

  其中StratifiedShuffleSplit对象的split方法会产生生成器(generator)。在python3中，无法使用generator.next()，而是使用next(generator)。

sss = StratifiedShuffleSplit(n_splits=5, random_state=0, test_size=1/6)
train_index, valid_index = next(sss.split(x_train, y_train))
x_valid, y_valid = x_train[valid_index], y_train[valid_index]
x_train, y_train = x_train[train_index], y_train[train_index]
print(x_train.shape, x_valid.shape, x_test.shape)
1
2
3
4
5

注意正确的结果应该是：
(50000, 28, 28, 1) (10000, 28, 28, 1) (10000, 28, 28, 1)，如果第一个一元组中的数不为50000，则需要再跑几次。

2. 使用TPU

2.1 查看环境是否为Colab TPU

import pprint

use_tpu = True #@param {type:"boolean"}

if use_tpu:
    assert 'COLAB_TPU_ADDR' in os.environ, 'Missing TPU; did you request a TPU in Notebook Settings?'

if 'COLAB_TPU_ADDR' in os.environ:
  TF_MASTER = 'grpc://{}'.format(os.environ['COLAB_TPU_ADDR'])
else:
  TF_MASTER=''

with tf.Session(TF_MASTER) as session:
  print ('List of devices:')
  pprint.pprint(session.list_devices())
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

2.2 使用GPU

resolver = tf.contrib.cluster_resolver.TPUClusterResolver(TF_MASTER)
tf.contrib.distribute.initialize_tpu_system(resolver)
strategy = tf.contrib.distribute.TPUStrategy(resolver)
1
2
3

3. 定义模型

def get_model():
  model = tf.keras.models.Sequential()
  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
  model.add(tf.keras.layers.Conv2D(64, (5, 5), padding='same', activation='elu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
  model.add(tf.keras.layers.Dropout(0.25))

  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
  model.add(tf.keras.layers.Conv2D(128, (5, 5), padding='same', activation='elu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
  model.add(tf.keras.layers.Dropout(0.25))

  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))
  model.add(tf.keras.layers.Conv2D(256, (5, 5), padding='same', activation='elu'))
  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))
  model.add(tf.keras.layers.Dropout(0.25))

  model.add(tf.keras.layers.Flatten())
  model.add(tf.keras.layers.Dense(256))
  model.add(tf.keras.layers.Activation('elu'))
  model.add(tf.keras.layers.Dropout(0.5))
  model.add(tf.keras.layers.Dense(10))
  model.add(tf.keras.layers.Activation('softmax'))
  model.summary()
  
  return model
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

4. 训练模型

with strategy.scope():
  model = get_model()
  model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3), 
                loss=tf.keras.losses.sparse_categorical_crossentropy,
                metrics=['sparse_categorical_accuracy'])
1
2
3
4
5

  其中batch_size必须为8的倍数，因为num_workders为8（在哪里设置不清楚，有知道的请留言）。steps_per_epoch也是必填项，为ceil(num_samples / batch_size)。

model.fit(
  x_train, y_train,
  epochs = 20,
  steps_per_epoch = 390,
  batch_size = 128
)
1
2
3
4
5
6

5. 评估模型

model.evaluate(
  x_test, y_test,batch_size=8
)
1
2
3