如何转换 tensorflow 模型以针对 TPU 进行优化？

Question

我一直在尝试使用 colab 提供的 TPU，因为它据说速度很快，但似乎无法做到。我正在使用张量流 2.4.1。我一直在尝试关注这个https://www.tensorflow.org/guide/tpu，但没有运气。这是代码 https://colab.research.google.com/drive/1GGtwBicZF0qtp57ioD7g0JdE1iBXL85J?usp=sharing

%tensorflow_version 2.x
from __future__ import absolute_import, division, print_function, unicode_literals
import os
from pathlib import Path
import tensorflow as tf
import numpy as np
import pandas as pd
import tensorflow_datasets as tfds
CSV_COLUMN_NAMES = ['SepalLength', 'SepalWidth', 'PetalLength', 'PetalWidth', 'Species']
SPECIES = ['Setosa', 'Versicolor', 'Virginica']
train_path = tf.keras.utils.get_file(
    "iris_training.csv", "https://storage.googleapis.com/download.tensorflow.org/data/iris_training.csv")
train = pd.read_csv(train_path, names=CSV_COLUMN_NAMES, header=0)
train_y = train.pop('Species')
nb_classes=3 # we have three types of flowers
X=np.array(train)
Y=np.eye(nb_classes)[np.array(train_y)]
clf = tf.keras.models.Sequential([
    tf.keras.layers.Dense(30, activation='relu'),
    tf.keras.layers.Dense(10, activation='relu'),
    tf.keras.layers.Dense(3, activation='softmax'),
])
clf.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
history = clf.fit(X,Y, batch_size=32,epochs=10, validation_split=0.1)
clf.save("numeric_values-model.h5")

这是我尝试转换它的尝试

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
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'))
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
with tf.device('/TPU:0'):
  c = tf.matmul(a, b)
print("c device: ", c.device)
print(c)
strategy = tf.distribute.TPUStrategy(resolver)
@tf.function
def matmul_fn(x, y):
  z = tf.matmul(x, y)
  return z

z = strategy.run(matmul_fn, args=(a, b))
print(z)
@tf.function
def matmul_fn(x, y):
  z = tf.matmul(x, y)
  return z

z = strategy.run(matmul_fn, args=(a, b))
print(z)
def create_model():
  return tf.keras.models.Sequential([
    tf.keras.layers.Dense(30, activation='relu'),
    tf.keras.layers.Dense(10, activation='relu'),
    tf.keras.layers.Dense(3, activation='softmax')])
def get_dataset(batch_size, is_training=True):
  split = 'train' if is_training else 'test'
  dataset, info = tfds.load(name='mnist', split=split, with_info=True,
                            as_supervised=True, try_gcs=True)


  # Only shuffle and repeat the dataset in training. The advantage to have a
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so users don't need to think about scaling the gradients
  # based on the actual batch size.
  if is_training:
    dataset = dataset.shuffle(10000)
    dataset = dataset.repeat()

  dataset = dataset.batch(batch_size)

  return dataset
with strategy.scope():
  model = create_model()
  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=['sparse_categorical_accuracy'])

batch_size = 200
steps_per_epoch = 60000 // batch_size
validation_steps = 10000 // batch_size

train_dataset = get_dataset(batch_size, is_training=True)
test_dataset = get_dataset(batch_size, is_training=False)

model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=steps_per_epoch,
          validation_data=test_dataset, 
          validation_steps=validation_steps)

Answer 1

您提到的代码 sn-p 失败并显示此错误消息：

  (0) Invalid argument: {{function_node __inference_train_function_10150}} Compilation failure: Incompatible shapes: [25,1] vs. [25,28,28]

意味着您的数据和输入之间存在形状不匹配。

这里最简单的解决方案是参考the official TensorFlow model garden MNIST model（它还使用了使用 TFDS 的 MNIST 上的简单模型）来找出您的代码的不同之处。