tensorflow - 关于如何在 TPU 云上运行自定义 keras 代码

Question

我编写的代码可以在 GPU 上运行，但实验的周转时间很长。我想移植这段代码，以便可以在 TPU 上运行它。我怎么能这样做？这就是我所拥有的。

Bunch of datloading stuff
!!!!!
!!!!!

effnet = efn.EfficientNetB5(weights=None, include_top=False)
effnet.load_weights(eff_weights_path)

# Replace all Batch Normalization layers by Group Normalization layers

model = Sequential()
model.add(effnet)
model.add(GlobalAveragePooling2D())
model.add(Dropout(0.5))
model.add(Dense(8,name ='elu', activation=elu))
model.load_weights('saved_models/wieghts_ef5.h5',by_name = True)
model.compile(loss='mse', optimizer='adam', metrics=['mse', 'acc'])
print(model.summary())        

model_json = model.to_json()
with open("./saved_models/model_ef5_fn.json", "w") as json_file:
    json_file.write(model_json)

BATCH_SIZE = 2
IMG_WIDTH  = 456
IMG_HEIGHT = 456
# Add Image augmentation to our generator
train_datagen = ImageDataGenerator(rotation_range=360,
                                   horizontal_flip=True,
                                   vertical_flip=True,
                                   validation_split=0.15,
                                   preprocessing_function=preprocess_image, 
                                   rescale=1 / 128.)

# Use the dataframe to define train and validation generators
train_generator = train_datagen.flow_from_dataframe(df, 
                                                    x_col='pic_id', 
                                                    y_col=Labels_list,
                                                    directory = train_images_path,
                                                    target_size=(IMG_WIDTH, IMG_HEIGHT),
                                                    batch_size=BATCH_SIZE,
                                                    class_mode='other', 
                                                    subset='training')

val_generator = train_datagen.flow_from_dataframe(df, 
                                                  x_col='pic_id', 
                                                  y_col=Labels_list,
                                                  directory = train_images_path,
                                                  target_size=(IMG_WIDTH, IMG_HEIGHT),
                                                  batch_size=BATCH_SIZE,
                                                  class_mode='other',
                                                  subset='validation')
print('Data was loaded')


# For tracking Quadratic Weighted Kappa score
kappa_metrics = Metrics()
# Monitor MSE to avoid overfitting and save best model
es = EarlyStopping(monitor='val_loss', mode='auto', verbose=1, patience=15)
rlr = ReduceLROnPlateau(monitor='val_loss', 
                        factor=0.5, 
                        patience=4, 
                        verbose=1, 
                        mode='auto', 
                        min_delta=0.0001)

# Begin training
model.fit_generator(train_generator,
                    # steps_per_epoch=60,
                    steps_per_epoch=train_generator.samples // BATCH_SIZE,                    
                    epochs=30,
                    validation_data=val_generator,
                    # validation_steps = 30,
                    validation_steps = val_generator.samples // BATCH_SIZE,                    
                    callbacks=[kappa_metrics, es, rlr])

Answer 1

正如聊天中所讨论的，这里是在 TPU 上训练模型的基本设计:

import os
import tensorflow as tf
from tensorflow.keras import *

!pip install -U git+https://github.com/qubvel/efficientnet
import efficientnet.tfkeras as efn

backend.clear_session()

if "COLAB_TPU_ADDR" not in os.environ:
    raise Exception("TPU not found")

tpu_addr = "grpc://" + os.environ["COLAB_TPU_ADDR"]
print("TPU address:", tpu_addr)

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu = tpu_addr)
tf.config.experimental_connect_to_host(resolver.master())
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

def get_model():
  effnet = efn.EfficientNetB5(weights = None, include_top = False)

  # Replace all Batch Normalization layers by Group Normalization layers
  for i, layer in enumerate(effnet.layers):
      if "batch_normalization" in layer.name:
          effnet.layers[i] = GroupNormalization(groups = 32, axis = -1, epsilon = 1e-5)

  model = models.Sequential()
  model.add(effnet)
  # add other layers
  return model

with strategy.scope():
  model = get_model()
  model.compile(optimizer = optimizers.Adagrad(learning_rate = 0.1), 
    loss = "sparse_categorical_crossentropy",
    metrics = ["sparse_categorical_crossentropy"]
  )

  #model.load_weights("/path/to/imagenet/weights")

print(model.summary())

# model.fit(...)
model.save_weights("weights.h5")

cpu_model = get_model()
cpu_model.load_weights("weights.h5")
#cpu_model.predict(...)