如果想尝试使用Google Colab上的TPU来训练模型，也是非常方便，仅需添加6行代码。
在Colab笔记本中：修改->笔记本设置->硬件加速器 中选择 TPU
注：以下代码只能在Colab 上才能正确执行。
https://colab.research.google.com/drive/1XCIhATyE1R7lq6uwFlYlRsUr5d9_-r1s
%tensorflow_version 2.x import tensorflow as tf print(tf.__version__) from tensorflow.keras import *
一，准备数据
MAX_LEN = 300 BATCH_SIZE = 32 (x_train,y_train),(x_test,y_test) = datasets.reuters.load_data() x_train = preprocessing.sequence.pad_sequences(x_train,maxlen=MAX_LEN) x_test = preprocessing.sequence.pad_sequences(x_test,1)">MAX_LEN) MAX_WORDS = x_train.max()+1 CAT_NUM = y_train.max()+1 ds_train = tf.data.Dataset.from_tensor_slices((x_train,y_train)) \ .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \ .prefetch(tf.data.experimental.AUTOTUNE).cache() ds_test = tf.data.Dataset.from_tensor_slices((x_test,y_test)) \ .shuffle(buffer_size = 1000).batch(BATCH_SIZE) \ .prefetch(tf.data.experimental.AUTOTUNE).cache()
二，定义模型
tf.keras.backend.clear_session() def create_model(): model = models.Sequential() model.add(layers.Embedding(MAX_WORDS,7,input_length=MAX_LEN)) model.add(layers.Conv1D(filters = 64,kernel_size = 5,activation = "relu")) model.add(layers.MaxPool1D(2)) model.add(layers.Conv1D(filters = 32,kernel_size = 3,1)">)) model.add(layers.Flatten()) model.add(layers.Dense(CAT_NUM,activation = softmax)) return(model) compile_model(model): model.compile(optimizer=optimizers.Nadam(),loss=losses.SparseCategoricalCrossentropy(from_logits=True),metrics=[metrics.SparseCategoricalAccuracy(),metrics.SparseTopKCategoricalAccuracy(5)]) return(model)
三，训练模型
# 增加以下6行代码 os resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ[COLAB_TPU_ADDR']) tf.config.experimental_connect_to_cluster(resolver) tf.tpu.experimental.initialize_tpu_system(resolver) strategy = tf.distribute.experimental.TPUStrategy(resolver) with strategy.scope(): model = create_model() model.summary() model = compile_model(model)
INFO:tensorflow:Initializing the TPU system: grpc://10.62.22.122:8470 INFO:tensorflow:Initializing the TPU system: grpc://10.62.22.122:8470 INFO:tensorflow:Clearing out eager caches INFO:tensorflow:Clearing out eager caches INFO:tensorflow:Finished initializing TPU system. INFO:tensorflow:Finished initializing TPU system. INFO:tensorflow:Found TPU system: INFO:tensorflow:Found TPU system: INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:cpu:0,cpu,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_cpu:0,XLA_cpu,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0,TPU,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1,0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6device:TPU_SYSTEM:0,TPU_SYSTEM,0) Model: sequential" _________________________________________________________________ Layer (type) Output Shape Param # ================================================================= embedding (Embedding) (None,300,7) 216874 conv1d (Conv1D) (None,296,64) 2304 max_pooling1d (MaxPooling1D) (None,148,64) 0 conv1d_1 (Conv1D) (None,146,32) 6176 max_pooling1d_1 (MaxPooling1 (None,73,32) 0 flatten (Flatten) (None,2336) 0 dense (Dense) (None,46) 107502 ================================================================= Total params: 332,856 Trainable params: 332,1)"> Non-trainable params: 0 _________________________________________________________________
history = model.fit(ds_train,validation_data = ds_test,epochs = 10)
前面的都没问题，最后运行上面这句话时colab崩溃了，colab自动重启，不知道是什么原因，下面是原书中的结果：
Train for 281 steps,validate for 71 steps Epoch 1/10 281/281 [==============================] - 12s 43ms/step - loss: 3.4466 - sparse_categorical_accuracy: 0.4332 - sparse_top_k_categorical_accuracy: 0.7180 - val_loss: 3.3179 - val_sparse_categorical_accuracy: 0.5352 - val_sparse_top_k_categorical_accuracy: 0.7195 Epoch 2/10 281/281 [==============================] - 6s 20ms/step - loss: 3.3251 - sparse_categorical_accuracy: 0.5405 - sparse_top_k_categorical_accuracy: 0.7302 - val_loss: 3.3082 - val_sparse_categorical_accuracy: 0.5463 - val_sparse_top_k_categorical_accuracy: 0.7235 Epoch 3/10 281/281 [==============================] - 6s 20ms/step - loss: 3.2961 - sparse_categorical_accuracy: 0.5729 - sparse_top_k_categorical_accuracy: 0.7280 - val_loss: 3.3026 - val_sparse_categorical_accuracy: 0.5499 - val_sparse_top_k_categorical_accuracy: 0.7217 Epoch 4/10 281/281 [==============================] - 5s 19ms/step - loss: 3.2751 - sparse_categorical_accuracy: 0.5924 - sparse_top_k_categorical_accuracy: 0.7276 - val_loss: 3.2957 - val_sparse_categorical_accuracy: 0.5543 - val_sparse_top_k_categorical_accuracy: 0.7217 Epoch 5/10 281/281 [==============================] - 5s 19ms/step - loss: 3.2655 - sparse_categorical_accuracy: 0.6008 - sparse_top_k_categorical_accuracy: 0.7290 - val_loss: 3.3022 - val_sparse_categorical_accuracy: 0.5490 - val_sparse_top_k_categorical_accuracy: 0.7231 Epoch 6/10 281/281 [==============================] - 5s 19ms/step - loss: 3.2616 - sparse_categorical_accuracy: 0.6041 - sparse_top_k_categorical_accuracy: 0.7295 - val_loss: 3.3015 - val_sparse_categorical_accuracy: 0.5503 - val_sparse_top_k_categorical_accuracy: 0.7235 Epoch 7/10 281/281 [==============================] - 6s 21ms/step - loss: 3.2595 - sparse_categorical_accuracy: 0.6059 - sparse_top_k_categorical_accuracy: 0.7322 - val_loss: 3.3064 - val_sparse_categorical_accuracy: 0.5454 - val_sparse_top_k_categorical_accuracy: 0.7266 Epoch 8/10 281/281 [==============================] - 6s 21ms/step - loss: 3.2591 - sparse_categorical_accuracy: 0.6063 - sparse_top_k_categorical_accuracy: 0.7327 - val_loss: 3.3025 - val_sparse_categorical_accuracy: 0.5481 - val_sparse_top_k_categorical_accuracy: 0.7231 Epoch 9/10 281/281 [==============================] - 5s 19ms/step - loss: 3.2588 - sparse_categorical_accuracy: 0.6062 - sparse_top_k_categorical_accuracy: 0.7332 - val_loss: 3.2992 - val_sparse_categorical_accuracy: 0.5521 - val_sparse_top_k_categorical_accuracy: 0.7257 Epoch 10/10 281/281 [==============================] - 5s 18ms/step - loss: 3.2577 - sparse_categorical_accuracy: 0.6073 - sparse_top_k_categorical_accuracy: 0.7363 - val_loss: 3.2981 - val_sparse_categorical_accuracy: 0.5516 - val_sparse_top_k_categorical_accuracy: 0.7306 cpu times: user 18.9 s,sys: 3.86 s,total: 22.7 s Wall time: 1min 1s
参考：
开源电子书地址：https://lyhue1991.github.io/eat_tensorflow2_in_30_days/
GitHub 项目地址：https://github.com/lyhue1991/eat_tensorflow2_in_30_days
原文链接：/tensorflow/991529.html