##tf2.4.0
from tensorflow.python.keras.callbacks import EarlyStopping
from tensorflow.python.keras.layers import Embedding, SpatialDropout1D, LSTM, Dense
from tensorflow.python.keras.models import Sequential
import tensorflow as tf
import os
##下面6行为GPU设置，若用GPU，则用这6行，其他的就不用了了
# gpus = tf.config.list_physical_devices("GPU")
# print(gpus)
# if gpus:
# gpu0 = gpus[0] # 如果有多个GPU，仅使用第0个GPU
# tf.config.experimental.set_memory_growth(gpu0, True) # 设置GPU显存用量按需使用
# tf.config.set_visible_devices([gpu0], "GPU")
#加载数据
with open('w2v1000.pkl', 'rb') as f:
dict = pickle.load(f)
X = dict['X']
Y = dict['Y']
#TPU设置代码，以下全为TPU的设置，需要注释掉上面GPU设置的6行，所有代码原封不动照搬即可，只需要
#在模型编译print语句后把自己模型给替换即可。
tf.keras.backend.clear_session()
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))
#strategy = tf.distribute.experimental.TPUStrategy(resolver)#这句好像已经改成下面了
strategy = tf.distribute.TPUStrategy(resolver)
with strategy.scope():#模型编译
print("# Preparing model")#就是这条语句后替换自己的模型即可，其他不用动
model = Sequential()#我写的示例为LSTM模型
model.add(LSTM(128, dropout=0.2, recurrent_dropout=0.2, input_shape=(1600, 150), return_sequences=True))
model.add(Dense(32))
#model.add(Dense(2, activation='softmax'))
model.add(Dense(1, activation='sigmoid'))
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['acc'])
print(" ", model.summary())
#模型训练，这条语句就按正常写即可，与TPU无关了
model.fit(X, Y, epochs=10, batch_size=10, validation_split=0.2,workers=4,
use_multiprocessing=True,callbacks=[EarlyStopping(monitor='loss', patience=7, min_delta=0.0001)])
#模型保存
model.save('w2vmodel1000.h5')
以下是Colab测试结果：
可以看到设备类型为TPU，速度十分快，同样模型用GPU（Tesla T4）一个Epoch需要800多秒，快了20到30倍，这还是bitchsize设置的比较小，所以TPU yyds！