Tensorflow：从任意长度的复杂张量中提取连续的补丁

Question

I'm trying to figure out how to extract sequential patches from a complex valued tensor where the length is variable. The extraction is being performed as part of a tf.data pipeline.

If the tensor were not complex, I'd use tf.image.extract_image_patches as in this answer .

However, that function does not work with complex tensors. I have tried the following technique, but it fails because the length of the tensor is unknown.

def extract_sequential_patches(image):
    image_length = tf.shape(image)[0]
    num_patches = image_length // (128 // 4)
    patches = []
    for i in range(num_patches):
        start = i * 128
        end = start + 128
        patches.append(image[start:end, ...])
    return tf.stack(patches)

However I get the error:

InaccessibleTensorError: The tensor 'Tensor("strided_slice:0", shape=(None, 512, 2), dtype=complex64)' cannot be accessed here: it is defined in another function or code block. Use return values, explicit Python locals or TensorFlow collections to access it. Defined in: FuncGraph(name=while_body_2100, id=140313967335120)

I have tried liberal decoration with @tf.function

Answer 1

I think you'll need to adjust the calculation of the indices to make sure they don't go out of bounds, but leaving that detail aside, your code is almost what tf.function expects, except for the use of a Python list; you need to use TensorArray instead.

Something like this should work (the index calculations might not be entirely right):

@tf.function
def extract_sequential_patches(image, size, stride):
    image_length = tf.shape(image)[0]
    num_patches = (image_length - size) // stride + 1
    patches = tf.TensorArray(image.dtype, size=num_patches)
    for i in range(num_patches):
        start = i * stride
        end = start + size
        patches = patches.write(i, image[start:end, ...])
    return patches.stack()

You can find more details on why Python lists don't currently work in the autograph reference docs .

That said, it might be faster to use the real/imag trick, if the kernel of extract_image_patches is optimized. I recommend testing both approaches.