TPU使用時のエラー：InvalidArgumentError: Unable to parse tensor proto

前提・実現したいこと

tensorflowで機械学習の画像分類を試みているものです

ImageNet2012をTensorFlow Datasetsに格納して、TPUの機械学習にかけようとしたのですがmodel.fitの箇所で下記エラーが出てしまいました。

InvalidArgumentError: Unable to parse tensor proto

解決方法があればご教示いただけますと幸いです。

発生している問題・エラーメッセージ

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-11-ab2ae0c92b62> in <module>()
     49                         epochs = 1000,
     50                         validation_data = ds_valid_,
---> 51                         shuffle = True,)

6 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
   1145           use_multiprocessing=use_multiprocessing,
   1146           model=self,
-> 1147           steps_per_execution=self._steps_per_execution)
   1148 
   1149       # Container that configures and calls `tf.keras.Callback`s.

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/data_adapter.py in get_data_handler(*args, **kwargs)
   1362   if getattr(kwargs["model"], "_cluster_coordinator", None):
   1363     return _ClusterCoordinatorDataHandler(*args, **kwargs)
-> 1364   return DataHandler(*args, **kwargs)
   1365 
   1366 

/usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/data_adapter.py in __init__(self, x, y, sample_weight, batch_size, steps_per_epoch, initial_epoch, epochs, shuffle, class_weight, max_queue_size, workers, use_multiprocessing, model, steps_per_execution, distribute)
   1148     else:
   1149       self._steps_per_execution = steps_per_execution
-> 1150       self._steps_per_execution_value = steps_per_execution.numpy().item()
   1151 
   1152     adapter_cls = select_data_adapter(x, y)

/usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/values.py in numpy(self)
    772   def numpy(self):
    773     if context.executing_eagerly():
--> 774       return self.read_value().numpy()
    775     else:
    776       raise NotImplementedError(

/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py in numpy(self)
   1092     """
   1093     # TODO(slebedev): Consider avoiding a copy for non-CPU or remote tensors.
-> 1094     maybe_arr = self._numpy()  # pylint: disable=protected-access
   1095     return maybe_arr.copy() if isinstance(maybe_arr, np.ndarray) else maybe_arr
   1096 

/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py in _numpy(self)
   1060       return self._numpy_internal()
   1061     except core._NotOkStatusException as e:  # pylint: disable=protected-access
-> 1062       six.raise_from(core._status_to_exception(e.code, e.message), None)  # pylint: disable=protected-access
   1063 
   1064   @property

/usr/local/lib/python3.7/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Unable to parse tensor proto

該当のソースコード

tensorflow dataset作成部分

Python
1import os
2import tensorflow_datasets as tfds
3
4write_dir = './ImageNet'
5
6# Construct a tf.data.Dataset
7download_config = tfds.download.DownloadConfig(extract_dir=os.path.join(write_dir, 'extracted'),
8                                               manual_dir='/content/drive/XXX/ILSVRC2012'
9                                               )
10download_and_prepare_kwargs = {'download_dir': os.path.join(write_dir, 'downloaded'),
11                               'download_config': download_config,
12                               }
13ds_train, ds_valid = tfds.load('imagenet2012_subset',
14                               data_dir=os.path.join(write_dir, 'data'),   
15                               split=['train', 'validation'], 
16                               shuffle_files=False, 
17                               download=True, 
18                               as_supervised=True,
19                               download_and_prepare_kwargs=download_and_prepare_kwargs)

エラーが出現した部分
TPU操作コードはリンクを参照

python
1from tensorflow.keras.callbacks import EarlyStopping, LearningRateScheduler, ModelCheckpoint
2from tensorflow.keras.datasets import cifar10
3from tensorflow.keras.layers import Activation, BatchNormalization, Conv2D, Dense, Flatten, Input, GlobalAveragePooling2D
4from tensorflow.keras.models import Model
5from tensorflow.keras.preprocessing.image import ImageDataGenerator
6from tensorflow.keras.utils import to_categorical
7from tensorflow.keras.applications.vgg16 import VGG16
8from tensorflow.keras.applications.inception_v3 import InceptionV3
9import tensorflow as tf
10import numpy as np
11import matplotlib.pyplot as plt
12%matplotlib inline
13
14#TPUの準備
15tpu_grpc_url = "grpc://" + os.environ["COLAB_TPU_ADDR"]
16tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
17tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
18tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)
19strategy = tf.distribute.experimental.TPUStrategy(tpu_cluster_resolver)
20
21
22with strategy.scope():
23    model = VGG16(input_shape = (224, 224, 3), )
24
25    # モデルのコンパイル
26    model.compile(optimizer='adam', loss = 'categorical_crossentropy', metrics= ['accuracy'])
27
28    # 画像を所定のサイズに変換
29    def resize_with_crop(image, label):
30        i = image
31        i = tf.image.resize_with_crop_or_pad(i, 224, 224)
32        i = tf.expand_dims(i,0)
33        l = tf.one_hot(label, 1000)
34        return (i, l)
35    ds_train_ = ds_train.map(resize_with_crop)
36    ds_valid_ = ds_valid.map(resize_with_crop)
37
38    history = model.fit(ds_train_,
39                        batch_size = 32,
40                        steps_per_epoch = 64,
41                        epochs = 1000,
42                        validation_data = ds_valid_,
43                        shuffle = True,)              #ここでエラー