TPU 上的张量切片

编程问答 2022-05-11

问题描述

我想在 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,:]
    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,tf.constant(lead_length,dtype=tf.int64))

    return rangeds.map(slice_,tf.data.experimental.AUTOTUNE)

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

输出 1 为 (432001,7) (432001,7,3)

输出 2 为 (None,7) (3,)

但是，我认为 OUTPUT2 应该是 (32768,)。

事实上，模型然后抱怨（只是来自一层的例子，还有更多，这是来自 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个元素是哪里来的...

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

解决方法

似乎存在使用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,:]
    y = tf.squeeze(tf.slice(annotations,[i+lead_length,[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,lead_length):
    X = data[i:i+lead_length,:]
    y = annotations[i+lead_length,:]
    X.set_shape((lead_length,7))
    return X,y

我认为在您的情况下，让 tf.slice 完成这项工作更干净。

tensorflowtensorflowtensorflowtf.kerastpu