TPU 上的张量切片

Question

我想在TPU(Google Cloud TPU)上运行一个模型。我试图减少到最低限度。我省略了模型代码，因为它不相关，我的问题发生得更早。

这是主要的python文件：

import tensorflow as tf
import os
from Model import Model
from DataGeneratorTPU import load_dataset
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=os.environ['TPU_NAME'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)
with strategy.scope():
model = Model(32768,7)
model.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])
dg = load_dataset('gs://bucket/data.tf','gs://bucket/annotations.tf',32768).batch(32,drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)
model.fit(dg,epochs=10,verbose=1)
model.save('test')

这是 DataGeneratorTPU.py：

import tensorflow as tf
def slice(i,data,annotations,lead_length):
X = data[i:i+lead_length,:]
y = annotations[i+lead_length,0,:]
print(X.shape,y.shape) #OUTPUT2
return X,y
def load_dataset(filename_data,filename_annotations,lead_length,step_size=1):
data = tf.io.parse_tensor(tf.io.read_file(filename_data), tf.float32)
annotations = tf.io.parse_tensor(tf.io.read_file(filename_annotations), tf.int32)
print(data.shape,annotations.shape) #OUTPUT1
rangeds = tf.data.Dataset.range(0,data.shape[0]-lead_length,step_size)
def slice_(i):
return slice(i,data,annotations,tf.constant(lead_length,dtype=tf.int64))
return rangeds.map(slice_, tf.data.experimental.AUTOTUNE)

您可能注意到，我用 OUTPUT1 和 OUTPUT2 标记了两个print语句，因此我可以告诉您输出是什么：

OUTPUT1(432001, 7) (432001, 7, 3)

OUTPUT2(None, 7) (3, )

但是，我认为OUTPUT2应该(32768, 7) (3, ).

事实上，模型随后抱怨(只是一层的例子，不止这些，这是来自 conv1d 层)：

(0) Invalid argument: {{function_node __inference_train_function_33579}} Compilation failure: Dynamic Spatial Convolution is not supported: lhs shape is f32[4,1,<=32774,7]     
[[{{node Model/conv1d/conv1d}}]]
TPU compilation failed
[[tpu_compile_succeeded_assert/_12623170171032432447/_5]] 
[[tpu_compile_succeeded_assert/_12623170171032432447/_5/_303]]

它抱怨我们正在谈论的维度(我在映射函数中打印)是动态的，而不是固定在 32768。但是，它应该是静态的，因为我对 32768 的切片使用恒定宽度，我什至确保范围不会查看最后可能出错的 32768 元素。似乎只是为了能够估计这比 32774 小，我不知道 6 个额外的元素是从哪里来的......

我做错了什么？我怎样才能得到这个静态？

Answer 1

似乎有一种情况，使用tf.strided_slice(这是__getitem__方法调用的函数)丢失传递给它的张量的形状信息。我想这是因为切片非常灵活，并且允许传递"不可能"的切片大小(例如，大于数组大小的end索引)。这样做会导致最终数据集中出现形状可变的元素。该函数无法确保有关数组最终形状的任何信息，因此它默认为None。

您的案例非常简单，可以替换为对tf.slice的调用，通过询问切片的大小来保留形状信息。

将slice函数替换为以下内容：

def slice(i, data, annotations, lead_length):
X = tf.slice(data, [i,0], [lead_length, tf.shape(data)[1]])
# I also used slice for y for the sake of it, but its probably more readable to use
# y = annotations[i+lead_length,0,:]
y = tf.squeeze(tf.slice(annotations, [i+lead_length,0,0], [1, 1, tf.shape(annotations)[2]]))
return X, y

查看数据集形状可以得到：

>>> ds = rangeds.map(slice_, tf.data.experimental.AUTOTUNE)
>>> ds
<ParallelMapDataset shapes: ((32768, 7), (3,)), types: (tf.float32, tf.float32)>

---

另一种可能性是，如果你知道你可以保证形状是正确的(即，i+lead_length永远不会大于你的第一维的大小)，那么在你的张量上调用set_shape。如果不能，将导致难以调试运行时错误。

def slice(i,data,annotations,lead_length):
X = data[i:i+lead_length,:]
y = annotations[i+lead_length,0,:]
X.set_shape((lead_length,7))
return X,y

我认为，就你而言，让tf.slice做这项工作更干净。