使用TPU加速训练,将Keras模型转TPU模型时报错,如图
关键代码如下
引用库:
%tensorflow_version 1.x
import json
import os
import numpy as np
import tensorflow as tf
from tensorflow.python.keras.applications import resnet
from tensorflow.python.keras import callbacks
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
转换TPU模型代码如下
# This address identifies the TPU we'll use when configuring TensorFlow.
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tf.logging.set_verbosity(tf.logging.INFO)
self.model = tf.contrib.tpu.keras_to_tpu_model(self.model, strategy=tf.contrib.tpu.TPUDistributionStrategy(tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))
self.model = resnet50.ResNet50(weights=None, input_shape=dataset.input_shape, classes=num_classes)