user5779223
user5779223

Reputation: 1490

Dynamically partition a 2d tensor into multiple tensors in tensorflow

Given a 2d tensor (matrix), I would like to partition it into several small ones with equal size. You can regard it as the preprocessing of the max pooling. For instance,

1 2 3 4 5 6 7 8
2 3 4 5 6 7 8 9
3 4 5 6 7 8 9 10 
4 5 6 7 8 9 10 11

Given the a dynamic desired_size of 2 * 4, the outputs should be:

1 2 3 4
2 3 4 5

5 6 7 8
6 7 8 9

3 4 5 6
4 5 6 7

7 8 9 10 
8 9 10 11

I have studied slice and gather for a while. But I still don't have idea how to do it. Could you tell me how to get that? Thanks in advance!

Upvotes: 1

Views: 1161

Answers (2)

Vijay Mariappan
Vijay Mariappan

Reputation: 17191

I tied with tf.split():

num_splits = 2
desired_size = (2, 4)
A = tf.constant(a)

C = tf.concat(tf.split(A, desired_size[0], 0),1)
D = tf.reshape(tf.concat(tf.split(C, num_splits*desired_size[0], 1), 0), (-1, desired_size[0], desired_size[1]))

#The result
[[[ 1  2  3  4]
[ 2  3  4  5]]

[[ 5  6  7  8]
[ 6  7  8  9]]

[[ 3  4  5  6]
[ 4  5  6  7]]

[[ 7  8  9 10]
[ 8  9 10 11]]]

# For num_splits = 4, desired_size = (2, 2) you get 
[[[ 1  2]
[ 2  3]]

[[ 3  4]
[ 4  5]]

[[ 5  6]
[ 6  7]]

[[ 7  8]
[ 8  9]]

[[ 3  4]
[ 4  5]]

[[ 5  6]
[ 6  7]]

[[ 7  8]
[ 8  9]]

[[ 9 10]
[10 11]]]

Upvotes: 1

P-Gn
P-Gn

Reputation: 24581

You could use tf.extract_image_patches, even though it turns out somewhat verbose:

import numpy as np
import tensorflow as tf

x = tf.constant(np.arange(8) + np.arange(1,5)[:,np.newaxis])
e = tf.extract_image_patches(x[tf.newaxis,:,:,tf.newaxis],
    [1, 2, 4, 1], [1, 2, 4, 1], [1, 1, 1, 1], padding='VALID')
e = tf.reshape(e, [-1, 2, 4])
sess = tf.InteractiveSession()
e.eval()

# returns
# array([[[ 1,  2,  3,  4],
#         [ 2,  3,  4,  5]],
#        [[ 5,  6,  7,  8],
#         [ 6,  7,  8,  9]],
#        [[ 3,  4,  5,  6],
#         [ 4,  5,  6,  7]],
#        [[ 7,  8,  9, 10],
#         [ 8,  9, 10, 11]]])

Upvotes: 1

Related Questions