tensorflow - TPU 上的张量切片
问题描述
我想在 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)
您可能会注意到,我print
用 OUTPUT1 和 OUTPUT2 标记了这两个语句,所以我可以告诉您输出是什么:
输出 1 是(432001, 7) (432001, 7, 3)
输出 2 是(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 个额外的元素是从哪里来的......
我究竟做错了什么?我怎样才能得到这个静态?
解决方案
似乎有一种情况,使用 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
工作更干净。
推荐阅读
- python - 如何使用一个数据框的索引在熊猫中重新索引另一个数据框
- python - 在 pd.read_csv(filename) 之后,我的行号是 NaN,为什么?
- python - 如何使用python计算直到字符串中的特定数字?
- python-3.x - 而不是在python中循环
- shell - 使用带有 curl 的密钥库
- python - 正则表达式排除后跟空格的单词
- javascript - 如何在距离一周结束的日子之后控制台记录一周中的某一天?
- reactjs - ReactJS 上下文 API 未正确更新刷新视图
- google-cloud-platform - 无论如何,是否可以根据计划而不是 GCP(谷歌云平台)中 dataproc 的策略自动扩展
- typescript - 符号是什么意思| 在打字稿中?