Best way to mimic PyTorch sliced assignment with Keras/Tensorflow - python

I am trying to mimic the operation done in PyTorch below:
vol = Variable(torch.FloatTensor(A, B*2, C, D, E).zero_()).cuda()
for i in range(C):
if i > 0 :
vol[:, :B, i, :,i:] = input0[:,:,:,i:]
vol[:, B:, i, :,i:] = input1[:,:,:,:-i]
else:
vol[:, :B, i, :,:] = input0
vol[:, B:, i, :,:] = input1
So far, I have tried using the following sliced assignment in TF and wrapping it in a Keras Lambda layer:
vol = tf.Variable(K.zeros((A, D, E, C, B*2)))
for i in range(C):
if i > 0:
vol[:, :, i:, i, :B].assign(input0[:,:,i:,:])
vol[:, :, i:, i, B:].assign(input1[:,:,:-i,:])
else:
vol[:, :, :, i, :B].assign(input0)
vol[:, :, :, i, B:].assign(input1)
return vol
I also tried vol = vol[...].assign(...).
This assigns the values to the vol variable correctly, which I can then convert to a tensor to use in the rest of my graph. However, the gradient of this operation is undefined in TF (LookupError: No gradient defined for operation 'strided_slice/_assign' (op type: StridedSliceAssign)), and the gradient doesn't get propagated to the previous layers that generate input0 and input1, while they do appear to get transferred in the PyTorch implementation. Is there a way to construct this same variable in TF such that the gradient is defined and my previous operations don't have a None gradient?

You need to construct the tensor "by hand". Assuming both input0 and input1 have shape (A, D, E, B), you can do something like this:
# Make the indexing mask with TensorFlow
in_shape = tf.shape(input0)
in_dims = 4
idx = tf.meshgrid(*[tf.range(in_shape[i]) for i in range(in_dims)], indexing='ij')[2]
idx = tf.expand_dims(idx, axis=3)
r = tf.range(C)[tf.newaxis, tf.newaxis, tf.newaxis, :, tf.newaxis]
mask = idx >= r
# If all dimensions are known at graph construction time, you can instead
# make the mask with NumPy like this to save graph computation time
idx = np.meshgrid(*[np.arange(d) for d in (A, D, E, B)], indexing='ij')[2]
idx = np.expand_dims(idx, 3)
r = np.arange(C)[np.newaxis, np.newaxis, np.newaxis, :, np.newaxis]
mask = idx >= r
# Make the tensor
input0_tile = tf.tile(tf.expand_dims(input0, 3), (1, 1, 1, C, 1))
input1_tile = tf.tile(tf.expand_dims(input1, 3), (1, 1, 1, C, 1))
zero_tile = tf.zeros_like(input0_tile)
vol0 = np.where(mask, input0_tile, zero_tile)
vol1 = np.where(mask, input1_tile, zero_tile)
vol = tf.concat([vol0, vol1], axis=-1)
Note that you need either the first or the second block followed by the third block, not the three blocks (see comments). The code builds a binary mask using a tf.meshgrid and a tf.range of indices, then uses tf.where to select values from the inputs or zeros.

A tf.Variable is sort of a primitive/basic type. You shouldn't want to gradients to propagate out of them.
What you want is to construct a node that outputs the 5 dimensional tensor like you want.
I would run a concatenate operation on the 4th dimension to build the tensor and use the result in place of the vol.
If you don't care about the gradients propagating to input0 and input1, then I would just build the tensor outside of tensorflow and use it as an initializer.

Related

Implement ConvND in Tensorflow

So I need a ND convolutional layer that also supports complex numbers. So I decided to code it myself.
I tested this code on numpy alone and it worked. Tested with several channels, 2D and 1D and complex. However, I have problems when I do it on TF.
This is my code so far:
def call(self, inputs):
with tf.name_scope("ComplexConvolution_" + str(self.layer_number)) as scope:
inputs = self._verify_inputs(inputs) # Check inputs are of expected shape and format
inputs = self.apply_padding(inputs) # Add zeros if needed
output_np = np.zeros( # I use np because tf does not support the assigment
(inputs.shape[0],) + # Per each image
self.output_size, # Image out size
dtype=self.input_dtype # To support complex numbers
)
img_index = 0
for image in inputs:
for filter_index in range(self.filters):
for i in range(int(np.prod(self.output_size[:-1]))): # for each element in the output
index = np.unravel_index(i, self.output_size[:-1])
start_index = tuple([a * b for a, b in zip(index, self.stride_shape)])
end_index = tuple([a+b for a, b in zip(start_index, self.kernel_shape)])
# set_trace()
sector_slice = tuple(
[slice(start_index[ind], end_index[ind]) for ind in range(len(start_index))]
)
sector = image[sector_slice]
new_value = tf.reduce_sum(sector * self.kernels[filter_index]) + self.bias[filter_index]
# I use Tied Bias https://datascience.stackexchange.com/a/37748/75968
output_np[img_index][index][filter_index] = new_value # The complicated line
img_index += 1
output = apply_activation(self.activation, output_np)
return output
input_size is a tuple of shape (dim1, dim2, ..., dim3, channels). An 2D rgb conv for example will be (32, 32, 3) and inputs will have shape (None, 32, 32, 3).
The output size is calculated from an equation I found in this paper: A guide to convolution arithmetic for deep learning
out_list = []
for i in range(len(self.input_size) - 1): # -1 because the number of input channels is irrelevant
out_list.append(int(np.floor((self.input_size[i] + 2 * self.padding_shape[i] - self.kernel_shape[i]) / self.stride_shape[i]) + 1))
out_list.append(self.filters)
Basically, I use np.zeros because if I use tf.zeros I cannot assign the new_value and I get:
TypeError: 'Tensor' object does not support item assignment
However, in this current state I am getting:
NotImplementedError: Cannot convert a symbolic Tensor (placeholder_1:0) to a numpy array.
On that same assignment. I don't see an easy fix, I think I should change the strategy of the code completely.
In the end, I did it in a very inefficient way based in this comment, also commented here but at least it works:
new_value = tf.reduce_sum(sector * self.kernels[filter_index]) + self.bias[filter_index]
indices = (img_index,) + index + (filter_index,)
mask = tf.Variable(tf.fill(output_np.shape, 1))
mask = mask[indices].assign(0)
mask = tf.cast(mask, dtype=self.input_dtype)
output_np = array * mask + (1 - mask) * new_value
I say inefficient because I create a whole new array for each assignment. My code is taking ages to compute for the moment so I will keep looking for improvements and post here if I get something better.

How to use indices from tf.nn.top_k with tf.gather_nd?

I'm trying to use the indices returned from tf.nn.top_k to extract the values from a second tensor.
I've tried using numpy type indexing, as well as tf.gather_nd directly, but I noticed that the indexing is wrong.
# temp_attention_weights of shape [I, B, 1, J]
top_values, top_indices = tf.nn.top_k(temp_attention_weights, k=top_k)
# top_indices of shape [I, B, 1, top_k], base_encoder_transformed of shape [I, B, 1, J]
# I now want to extract from base_encoder_transformed top_indices
base_encoder_transformed = tf.gather_nd(base_encoder_transformed, indices=top_indices)
# base_encoder_transformed should be of shape [I, B, 1, top_k]
I noticed that top_indices is of the wrong format, but I can't seem to transform it to be used in tf.gather_nd, where the innermost dimension is used to index each corresponding element from base_encoder_transformed. Does anybody know a way to get top_indices into the right format?
top_indices will index only over the last axis, you need to add indices for the rest of axes too. That is easy with tf.meshgrid:
import tensorflow as tf
# Example input data
I = 4
B = 3
J = 5
top_k = 2
x = tf.reshape(tf.range(I * B * J), (I, B, 1, J)) % 7
# Top K
top_values, top_indices = tf.nn.top_k(x, k=top_k)
# Make indices for the rest of axes
ii, jj, kk, _ = tf.meshgrid(
tf.range(I),
tf.range(B),
tf.range(1),
tf.range(top_k),
indexing='ij')
# Stack complete index
index = tf.stack([ii, jj, kk, top_indices], axis=-1)
# Get the same values again
top_values_2 = tf.gather_nd(x, index)
# Test
with tf.Session() as sess:
v1, v2 = sess.run([top_values, top_values_2])
print((v1 == v2).all())
# True
I don't see a reason to use tf.gather_nd. There is an easier and faster (no need to use tf.meshgrid) solution using tf.gather with the batch_dims parameter.
import tensorflow as tf
# Example input data
I = 4
B = 3
J = 5
top_k = 2
x = tf.reshape(tf.range(I * B * J), (I, B, 1, J)) % 7
# Top K
top_values, top_indices = tf.nn.top_k(x, k=top_k)
#Gather indices along last axis
top_values_2 = tf.gather(x, top_indices, batch_dims = 3)
tf.reduce_all(top_values_2 == top_values).numpy()
#True
Note that batch_dims is 3 in this case because we want to gather from the last axis, and the rank of x is 4.

map function in Pytorch

Is there any map function in Pytorch? (something like map in python).
I need to map a 1xDxhxw tensor variable to a 1x(9D)xhxw tensor, to augment embedding of each pixel with its 8 neighbour embeddings. Is there any functionality in Pytorch that lets me do that efficiently?
I tried using map in Python this way:
n, d, h, w = embedding.size()
padder = nn.ReflectionPad2d(padding=1)
embedding = padder(embedding)
embedding = map(lambda i, j, M: M[:, :, i-1:i+2, j-1:j+2], range(1, h), range(1, w), embedding)
But it does not work for w > 2 and h > 2.
From your question, it is not clear what you are attempting to accomplish.
Note that full Python is supported in PyTorch, but what you are doing is creating a map object in your last line of code. The following should work for your purposes (? I'm guessing) though:
import torch
import torch.nn as nn
n, d, h, w = 20, 3, 32, 32
embedding = torch.randn(n, d, h, w)
padder = nn.ReflectionPad2d(padding=1)
embedding = padder(embedding)
new = [embedding[:,:, (i-1):(i+2), (j-1):(j+2)] for i, j in zip(range(1,h), range(1,w))]
Note however, that there are more elegant ways to chunk up a tensor (e.g. torch.chunk() or to operate on patches with convolutions (e.g. torch.nn.Conv2d)

Determining tensor shapes at time of graph creation in TensorFlow

I'm trying to write a chunk of reusable code that reads the shape of one tensor and then uses the resulting object to define the shape of other tensors. I have a choice of reading the dynamic shape of the tensor with tf.shape(tensor) or the static shape of the tensor with tensor.get_shape(). The toy example looks like this (with the two different strategies):
def my_function_strategy_1(x, y):
x_shape = tf.shape(x)
a = tf.reshape(y, x_shape)
b = tf.zeros(x_shape)
num_x_values = x_shape[0]
c = tf.reshape(y, [num_x_values, 4])
d = tf.zeros([num_x_values, 4])
return a, b, c, d
def my_function_strategy_2(x, y):
x_shape = x.get_shape()
a = tf.reshape(y, x_shape)
b = tf.zeros(x_shape)
num_x_values = x_shape[0]
c = tf.reshape(y, [num_x_values, 4])
d = tf.zeros([num_x_values, 4])
return a, b, c, d
I want to use this chunk of code in different graphs. Sometimes the shape of the input tensors will be known and sometimes they will be unknown:
graph_A = tf.Graph()
with graph_A.as_default():
x = tf.placeholder(tf.float32, [2, 4])
y = tf.placeholder(tf.float32, [8])
a, b, c, d = my_function(x, y)
with graph_B.as_default():
x = tf.placeholder(tf.float32)
y = tf.placeholder(tf.float32)
a, b, c, d = my_function(x, y)
The behavior I want is: (A) When the shapes of the input tensors are known (as in graph_A), I want TensorFlow to calculate all of the shapes in the graph at graph creation time (so it can efficiently allocate resources, etc.), and (B) When the shapes of the input tensors are unknown (as in graph_B), I want the TensorFlow to wait until runtime to calculate all of the shapes in the graph.
The strategy_1 version of the function almost does this. It achieves (B), but it doesn't quite achieve (A) because TensorFlow leaves the shape of some tensors unknown. For example, in the toy example above, the shapes of a, b, and c are calculated at graph creation time, but the shape of d is left unknown (even though d uses very similar operations). You can check this by printing a.get_shape(), b.get_shape(), etc.
Conversely, the strategy_2 version of the function achieves (A) for all tensors in the graph, but doesn't achieve (B) because TensorFlow (understandably) throws an exception when it tries to use the (unknown) static shape of the input tensor to shape other tensors.
Is there a way to achieve both (A) and (B) in a single function? How/why does the strategy_1 version work for most tensors in the graph, but not all?
You can carefully pick the elements of the shape that you know to have a "best of both worlds" result:
def my_get_shape(tensor):
if tensor.shape.ndims is None:
# Fully dynamic
return tf.shape(tensor)
if tensor.shape.is_fully_defined():
# Fully static
return tensor.shape
# Partially static
dyn_shape = tf.shape(tensor)
shape = []
for i, d in enumerate(tensor.shape):
shape.append(d.value if d.value is not None else dyn_shape[i])
return shape
def my_function(x, y):
x_shape = my_get_shape(x) # Or just tf.shape(x)! - see edit
a = tf.reshape(y, x_shape)
b = tf.zeros(x_shape)
num_x_values = x_shape[0]
c = tf.reshape(y, [num_x_values, 4])
d = tf.zeros([num_x_values, 4])
return a, b, c, d
# Fully static
with tf.Graph().as_default():
x = tf.placeholder(tf.float32, [2, 4])
y = tf.placeholder(tf.float32, [8])
a, b, c, d = my_function(x, y)
print('a:', a.shape, ', b:', b.shape, ', c:', c.shape, ', d:', d.shape)
# a: (2, 4) , b: (2, 4) , c: (2, 4) , d: (2, 4)
# Fully dynamic
with tf.Graph().as_default():
x = tf.placeholder(tf.float32)
y = tf.placeholder(tf.float32)
a, b, c, d = my_function(x, y)
print('a:', a.shape, ', b:', b.shape, ', c:', c.shape, ', d:', d.shape)
# a: <unknown> , b: <unknown> , c: (?, 4) , d: (?, 4)
# Partially static
with tf.Graph().as_default():
x = tf.placeholder(tf.float32, [None, 4])
y = tf.placeholder(tf.float32)
a, b, c, d = my_function(x, y)
print('a:', a.shape, ', b:', b.shape, ', c:', c.shape, ', d:', d.shape)
# a: (?, 4) , b: (?, 4) , c: (?, 4) , d: (?, 4)
EDIT:
Actually, replacing my_get_shape with tf.shape in the previous snippet works exacly the same. It seems that tf.shape should be the default (being careful not to cram the graph with it) unless you explicitly want to keep dimensions undefined.
I have investigated a bit, and I couldn't work the whole thing out completely. I don't know if this is useful, but here are some things I found out. Apparently TensorFlow has, at C++ level (it seems it used to be in Python before, but not anymore), a "shape inference" mechanism. If you look, for example, in tensorflow/core/ops/array_ops.cc) you will see that every operation declaration includes a .SetShapeFn at the end, which is a function that uses InferenceContext to try to guess the output shape of the operation. This class can, among other things , check whether values in a tensor are already known, which is true for example for tf.shape when the given tensor is static or for tf.fill (and related like tf.ones) with known values. The resolution of the shape inference algorithm is what is set as tensor shape in Python, and it can be called directly (although I don't see how it can be useful) through call_cpp_shape_fn:
from tensorflow.python.framework.common_shapes import call_cpp_shape_fn
with tf.Graph().as_default():
print(call_cpp_shape_fn(tf.reshape(tf.placeholder(tf.float32), tf.fill([2], 3)).op))
# Shows this:
# {
# 'shapes': [dim { size: 3 } dim { size: 3 }],
# 'handle_data': [None],
# 'inputs_needed': b'\x12\x01\x01'
# }
print(call_cpp_shape_fn(tf.reshape(tf.placeholder(tf.float32), (2 * tf.fill([2], 3))).op))
# Shows this:
# {
# 'shapes': [dim { size: -1 } dim { size: -1 }],
# 'handle_data': [None],
# 'inputs_needed': b'\x12\x01\x01'
# }
You can see that, while tf.fill([2], 3) was correctly inspected, TensorFlow didn't work out that 2 * tf.fill([2], 3) is [6, 6], presumably because statically keeping track of operations like multiplication, even if operands are known constants, was deemed too expensive.
What I haven't found out is where do ops declare that their values can be statically known, or where/how these values are retrieved exactly. It seems that, for example, for tf.shape, it is able to specifically pick known values and leave the rest as undefined.

How TensorArray and while_loop work together in tensorflow?

I am trying to produce a very easy example for combination of TensorArray and while_loop:
# 1000 sequence in the length of 100
matrix = tf.placeholder(tf.int32, shape=(100, 1000), name="input_matrix")
matrix_rows = tf.shape(matrix)[0]
ta = tf.TensorArray(tf.float32, size=matrix_rows)
ta = ta.unstack(matrix)
init_state = (0, ta)
condition = lambda i, _: i < n
body = lambda i, ta: (i + 1, ta.write(i,ta.read(i)*2))
# run the graph
with tf.Session() as sess:
(n, ta_final) = sess.run(tf.while_loop(condition, body, init_state),feed_dict={matrix: tf.ones(tf.float32, shape=(100,1000))})
print (ta_final.stack())
But I am getting the following error:
ValueError: Tensor("while/LoopCond:0", shape=(), dtype=bool) must be from the same graph as Tensor("Merge:0", shape=(), dtype=float32).
Anyone has on idea what is the problem?
There are several things in your code to point out. First, you don't need to unstack the matrix into the TensorArray to use it inside the loop, you can safely reference the matrix Tensor inside the body and index it using matrix[i] notation. Another issue is the different data type between your matrix (tf.int32) and the TensorArray (tf.float32), based on your code you're multiplying the matrix ints by 2 and writing the result into the array so it should be int32 as well. Finally, when you wish to read the final result of the loop, the correct operation is TensorArray.stack() which is what you need to run in your session.run call.
Here's a working example:
import numpy as np
import tensorflow as tf
# 1000 sequence in the length of 100
matrix = tf.placeholder(tf.int32, shape=(100, 1000), name="input_matrix")
matrix_rows = tf.shape(matrix)[0]
ta = tf.TensorArray(dtype=tf.int32, size=matrix_rows)
init_state = (0, ta)
condition = lambda i, _: i < matrix_rows
body = lambda i, ta: (i + 1, ta.write(i, matrix[i] * 2))
n, ta_final = tf.while_loop(condition, body, init_state)
# get the final result
ta_final_result = ta_final.stack()
# run the graph
with tf.Session() as sess:
# print the output of ta_final_result
print sess.run(ta_final_result, feed_dict={matrix: np.ones(shape=(100,1000), dtype=np.int32)})

Categories