##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设置，其实可以不写
# 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设置代码
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()
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())
#模型训练
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')
|