При работе на экземпляре TPU в Google Colab возникает InternalError: не удалось сериализовать сообщение.

В последних выпусках поддержка моделей Keras на TPU была значительно улучшена. Я пошел дальше и обновил ваш образец кода для TF 2.2. Большинство изменений - это простые переименования, и самое большое изменение состоит в том, что я настроил ваш набор входных данных с помощью _ 1_. Для достижения наилучших результатов на TPU мы всегда рекомендуем использовать tf.data.Dataset вместо использования numpy массивов напрямую с model.fit. Если у вас уже есть данные в numpy, вы можете создать набор данных с помощью tf.data.Dataset.from_tensor_slices((X_train, Y_train)) , хотя вы можете получить лучшие результаты, используя TFRecords. У меня нет доступа к вашему исходному набору данных, поэтому я пошел дальше и подставил вместо него случайные тензоры.

Вот обновленный код:

%tensorflow_version 2.x
import os
import tensorflow as tf
import numpy as np

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)

# define an encoder block
def define_encoder_block(layer_in, n_filters, batchnorm=True):
  # weight initialization
  init = tf.keras.initializers.RandomNormal(stddev=0.02)
  # add downsampling layer    
  g = tf.keras.layers.Conv2D(int(n_filters/2), (3,3), padding='same', kernel_initializer=init)(layer_in)
  g = tf.keras.layers.Conv2D(int(n_filters/2), (3,3), strides=(2,2), padding='same', kernel_initializer=init)(g)    
  g = tf.keras.layers.Conv2D(n_filters, (3,3), padding='same', kernel_initializer=init)(g)  
  # conditionally add batch normalization
  if batchnorm:
      g = tf.keras.layers.BatchNormalization()(g, training=True)
  # leaky relu activation
  g = tf.keras.activations.elu(g)
  return g


# define a decoder block
def decoder_block(layer_in, skip_in, n_filters, dropout=True):
  # weight initialization
  init = tf.keras.initializers.RandomNormal(stddev=0.02)
  # add upsampling layer
  g = tf.keras.layers.Conv2D(int(n_filters/2), (3,3),  padding='same', kernel_initializer=init)(layer_in)
  g = tf.keras.layers.Conv2D(int(n_filters/2), (3,3),  padding='same', kernel_initializer=init)(layer_in)    
  g = tf.keras.layers.Conv2DTranspose(n_filters, (3,3), strides=(2,2),  padding='same', kernel_initializer=init)(g)    
  # add batch normalization
  g = tf.keras.layers.BatchNormalization()(g, training=True)
  # conditionally add dropout
  if dropout:
      g = tf.keras.layers.Dropout(0.5)(g, training=True)
  # merge with skip connection
  g = tf.keras.layers.Concatenate()([g, skip_in])
  # relu activation
  g = tf.keras.activations.elu(g)
  return g

# define complete model
def define_generator(image_shape=(256,256,3)):
  # weight initialization
  init = tf.keras.initializers.RandomNormal(stddev=0.02)
  # image input
  in_image = tf.keras.layers.Input(shape=image_shape)
  # encoder model: C64-C128-C256-C512-C512-C512-C512-C512
  e1 = define_encoder_block(in_image, 64, batchnorm=False)
  e2 = define_encoder_block(e1, 128)
  e3 = define_encoder_block(e2, 256)
  e4 = define_encoder_block(e3, 512)
  e5 = define_encoder_block(e4, 512)
  e6 = define_encoder_block(e5, 512)
  e7 = define_encoder_block(e6, 512)
  # bottleneck, no batch norm and relu
  b = tf.keras.layers.Conv2D(512, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(e7)
  b = tf.keras.activations.elu(b)
  # decoder model: CD512-CD1024-CD1024-C1024-C1024-C512-C256-C128
  d1 = decoder_block(b, e7, 512)
  d2 = decoder_block(d1, e6, 512)
  d3 = decoder_block(d2, e5, 512)
  d4 = decoder_block(d3, e4, 512, dropout=False)
  d5 = decoder_block(d4, e3, 256, dropout=False)
  d6 = decoder_block(d5, e2, 128, dropout=False)
  d7 = decoder_block(d6, e1, 64, dropout=False)
  # output
  g = tf.keras.layers.Conv2DTranspose(3, (3,3), strides=(2,2), padding='same', kernel_initializer=init)(d7)
  out_image = tf.keras.activations.tanh(g)
  # define model
  model = tf.keras.models.Model(in_image, out_image)
  return model

# Values from original notebook
# shape = (11612,256,256,3) # this caused my notebook to OOM since it's huge
shape = (256,256,256,3)
batch_size = 8
epochs = 64

# Create fake random dataset
X_train = np.random.rand(*shape)
Y_train = np.random.rand(*shape)
dataset = (tf.data.Dataset.from_tensor_slices((X_train, Y_train))
    .repeat(epochs)
    .batch(batch_size, drop_remainder=True)
    .prefetch(16))

with strategy.scope():
  model = define_generator()
  adam = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5, beta_2=0.999)
  model.compile(optimizer=adam, loss='mean_absolute_error', metrics=['accuracy'])
  model.summary()

model.fit(dataset)