【tensorflow2.0】使用TPU训练模型

如果想尝试使用Google Colab上的TPU来训练模型，也是非常方便，仅需添加6行代码。

在Colab笔记本中：修改->笔记本设置->硬件加速器中选择 TPU

注：以下代码只能在Colab 上才能正确执行。

可通过以下colab链接测试效果《tf_TPU》：

https://colab.research.google.com/drive/1XCIhATyE1R7lq6uwFlYlRsUr5d9_-r1s

%tensorflow_version 2.x
@H_301_13@import@H_301_13@ tensorflow as tf
@H_301_13@print@H_301_13@(tf.__version__@H_301_13@)
@H_301_13@from@H_301_13@ tensorflow.keras import@H_301_13@ *

一，准备数据

MAX_LEN = 300
BATCH_SIZE @H_301_13@= 32
(x_train,y_train),(x_test,y_test) @H_301_13@= datasets.reuters.load_data()
x_train @H_301_13@= preprocessing.sequence.pad_sequences(x_train,maxlen=MAX_LEN)
x_test @H_301_13@= preprocessing.sequence.pad_sequences(x_test,1)">MAX_LEN)
 
MAX_WORDS @H_301_13@= x_train.max()+1
CAT_NUM @H_301_13@= y_train.max()+1
 
ds_train @H_301_13@= tf.data.Dataset.from_tensor_slices((x_train,y_train)) \
          .shuffle(buffer_size @H_301_13@= 1000).batch(BATCH_SIZE) \
          .prefetch(tf.data.experimental.AUTOTUNE).cache()
 
ds_test @H_301_13@= tf.data.Dataset.from_tensor_slices((x_test,y_test)) \
          .shuffle(buffer_size @H_301_13@= 1000).batch(BATCH_SIZE) \
          .prefetch(tf.data.experimental.AUTOTUNE).cache()@H_301_13@

二，定义模型

tf.keras.backend.clear_session()
@H_301_13@def@H_301_13@ create_model():
 
    model @H_301_13@= models.Sequential()
 
    model.add(layers.Embedding(MAX_WORDS,@H_301_13@7,input_length=MAX_LEN))
    model.add(layers.Conv1D(filters @H_301_13@= 64,kernel_size = 5,activation = "@H_301_13@relu@H_301_13@"@H_301_13@))
    model.add(layers.MaxPool1D(@H_301_13@2))
    model.add(layers.Conv1D(filters @H_301_13@= 32,kernel_size = 3,1)">))
    model.add(layers.Flatten())
    model.add(layers.Dense(CAT_NUM,activation @H_301_13@= softmax@H_301_13@))
    @H_301_13@return@H_301_13@(model)
 
@H_301_13@ compile_model(model):
    model.compile(optimizer@H_301_13@=optimizers.Nadam(),loss@H_301_13@=losses.SparseCategoricalCrossentropy(from_logits=True),metrics@H_301_13@=[metrics.SparseCategoricalAccuracy(),metrics.SparseTopKCategoricalAccuracy(5)]) 
    @H_301_13@return@H_301_13@(model)

三，训练模型

#@H_301_13@ 增加以下6行代码@H_301_13@
 os
resolver @H_301_13@= tf.distribute.cluster_resolver.TPUClusterResolver(tpu='@H_301_13@grpc://@H_301_13@'@H_301_13@ + os.environ[COLAB_TPU_ADDR@H_301_13@'@H_301_13@])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy @H_301_13@= tf.distribute.experimental.TPUStrategy(resolver)
with strategy.scope():
    model @H_301_13@= create_model()
    model.summary()
    model @H_301_13@= compile_model(model)

INFO:tensorflow:Initializing the TPU system: grpc://10.62.22.122:8470
INFO:tensorflow:Initializing the TPU system: grpc:@H_301_13@//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:@H_301_13@*** Num TPU Cores: 8
INFO:tensorflow:@H_301_13@*** Num TPU Cores: 8
INFO:tensorflow:@H_301_13@*** Num TPU Workers: 1
INFO:tensorflow:@H_301_13@*** Num TPU Cores Per Worker: 8
INFO:tensorflow:@H_301_13@*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:cpu:0,cpu,0)
INFO:tensorflow:@H_301_13@*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_cpu:0,XLA_cpu,0)
INFO:tensorflow:@H_301_13@*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0,TPU,0)
INFO:tensorflow:@H_301_13@*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1,0)
INFO:tensorflow:@H_301_13@*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2301_13@*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3301_13@*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4301_13@*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5301_13@*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6301_13@*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7device:TPU_SYSTEM:0,TPU_SYSTEM,0)
Model: @H_301_13@sequential@H_301_13@"@H_301_13@
_________________________________________________________________@H_301_13@
Layer (type)                 Output Shape              Param @H_301_13@#@H_301_13@   
=================================================================
embedding (Embedding)        (None,@H_301_13@300,7)            216874    

conv1d (Conv1D)              (None,@H_301_13@296,64)           2304      

max_pooling1d (MaxPooling1D) (None,@H_301_13@148,64)           0         
@H_301_13@
conv1d_1 (Conv1D)            (None,@H_301_13@146,32)           6176      

max_pooling1d_1 (MaxPooling1 (None,@H_301_13@73,32)            0         
@H_301_13@
flatten (Flatten)            (None,@H_301_13@2336)              0         
@H_301_13@
dense (Dense)                (None,@H_301_13@46)                107502    
=================================================================
Total params: @H_301_13@332,856
Trainable params: @H_301_13@332,1)">
Non@H_301_13@-trainable params: 0
@H_301_13@_________________________________________________________________@H_301_13@

history = model.fit(ds_train,validation_data = ds_test,epochs = 10)

前面的都没问题，最后运行上面这句话时colab崩溃了，colab自动重启，不知道是什么原因，下面是原书中的结果：

Train for@H_301_13@ 281 steps,validate for@H_301_13@ 71 steps
Epoch @H_301_13@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 @H_301_13@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 @H_301_13@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 @H_301_13@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 @H_301_13@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 @H_301_13@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 @H_301_13@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 @H_301_13@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 @H_301_13@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 @H_301_13@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 @H_301_13@18.9 s,sys: 3.86 s,total: 22.7 s
Wall time: 1min 1s@H_301_13@

参考：

开源电子书地址：https://lyhue1991.github.io/eat_tensorflow2_in_30_days/

GitHub 项目地址：https://github.com/lyhue1991/eat_tensorflow2_in_30_days

【tensorflow2.0】使用TPU训练模型

一，准备数据

二，定义模型

三，训练模型

猜你在找的Tensorflow相关文章