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)>

---

另一种可能性是在你的张量上调用set_shape如果你知道你可以保证形状是正确的（即i+lead_length永远不会比你的第一个尺寸大方面）。如果不能，将导致难以调试运行时错误。

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 做这项工作更干净。