def _prepare_gradient_func(
self, input_shape: ImageShape, input_layer: str,
output_layers: List[str]
) -> Tuple[Callable, GridShape, List[GridShape]]:
"""
Computes gradient function and additional parameters. Note
that the receptive field parameters like stride or size, do not
depend on input image shape. However, if the RF of original network
is bigger than input_shape this method will fail. Hence it is
recommended to increase the input shape.
:param input_shape: shape of the input image. Used in @model_func.
:param input_layer: name of the input layer.
:param output_layers: a list of names of the target feature map layers.
:returns
gradient_function: a function which returns gradient w.r.t. to
the input image
input_shape: a shape of the input image tensor
output_shape: a shapes of the output feature map tensors
"""
model = self._model_func(ImageShape(*input_shape))
if self.init_weights:
setup_model_weights(model)
gradient_function, input_shape, output_shapes = _define_receptive_field_func(
model, input_layer, output_layers)
return (
gradient_function,
GridShape(*input_shape),
[GridShape(*output_shape) for output_shape in output_shapes],
)
def _prepare_gradient_func(
self, input_shape: ImageShape, input_layer: str,
output_layer: str) -> Tuple[Callable, GridShape, GridShape]:
model = self.model_func(ImageShape(*input_shape))
if self.init_weights:
setup_model_weights(model)
gradient_function, input_shape, output_shape = \
_define_receptive_field_func(model, input_layer, output_layer)
return gradient_function, \
GridShape(*input_shape), \
GridShape(*output_shape)
def _prepare_gradient_func(
self, input_shape: ImageShape
) -> Tuple[Callable, GridShape, List[GridShape]]:
"""
Computes gradient function and additional parameters. Note
that the receptive field parameters like stride or size, do not
depend on input image shape. However, if the RF of original network
is bigger than input_shape this method will fail. Hence it is
recommended to increase the input shape.
:param input_shape: shape of the input image, which is feed to the
Pytorch module.
:returns
gradient_function: a function which returns gradient w.r.t. to
the input image.
input_shape: a shape of the input image tensor.
output_shapes: a list shapes of the output feature map tensors.
"""
input_shape = ImageShape(*input_shape)
input_shape = GridShape(
n=1,
c=input_shape.c,
h=input_shape.h,
w=input_shape.w,
)
gradient_function, input_shape, output_shapes = \
_define_receptive_field_func(self._model_func, input_shape)
return gradient_function, input_shape, output_shapes
def _get_gradient_from_grid_point(self,
gradient_function: Callable,
input_shape: GridShape,
output_shape: GridShape,
point: GridPoint,
intensity: float = 1.0) -> np.ndarray:
output_shape = output_shape._replace(n=1)
input_shape = input_shape._replace(n=1)
output_feature_map = np.zeros(shape=output_shape)
output_feature_map[:, point.x, point.y, 0] = intensity
receptive_field_grad = gradient_function(
[output_feature_map,
np.zeros(shape=input_shape), 0])[0]
return receptive_field_grad
def _prepare_gradient_func(
self,
input_shape: ImageShape,
input_layer: str,
output_layer: str
) -> Tuple[Callable, GridShape, GridShape]:
# this function will create default graph
_ = self.model_func(ImageShape(*input_shape))
default_graph = tf.get_default_graph()
# get graph tensors by names
input_tensor = default_graph \
.get_operation_by_name(input_layer).outputs[0]
output_tensor = default_graph \
.get_operation_by_name(output_layer).outputs[0]
# get their shapes
output_shape = _get_tensor_shape(output_tensor)
input_shape = _get_tensor_shape(input_tensor)
# define loss function
output_shape = (1, output_shape[1], output_shape[2], 1)
receptive_field_mask = tf.placeholder(
tf.float32, shape=output_shape, name='grid'
)
x = tf.reduce_mean(output_tensor, -1, keep_dims=True)
fake_loss = x * receptive_field_mask
fake_loss = tf.reduce_mean(fake_loss)
grads = tf.gradients(fake_loss, input_tensor)
# here we use Keras API to define gradient function which is simpler
# than native tf
gradient_function = keras.backend.function(
inputs=[receptive_field_mask, input_tensor, keras.backend.learning_phase()],
outputs=grads
)
return gradient_function, \
GridShape(*input_shape), \
GridShape(*output_shape)
def _define_receptive_field_func(model_fn: Callable[[], nn.Module],
input_shape: GridShape):
shape = [input_shape.n, input_shape.c, input_shape.h, input_shape.w]
input_image = torch.zeros(*shape)
model = model_fn()
_ = model(input_image)
feature_maps = model.__dict__.get("feature_maps", None)
if feature_maps is None:
raise ValueError("Field module.feature_maps is not defined. "
"Cannot compute receptive fields.")
if type(feature_maps) != list:
raise ValueError("Field module.feature_maps must be a list. "
"Cannot compute receptive fields.")
# compute feature maps output shapes
output_shapes = []
for feature_map in feature_maps:
output_shape = feature_map.size()
output_shape = GridShape(n=output_shape[0],
c=output_shape[1],
h=output_shape[2],
w=output_shape[3])
output_shapes.append(output_shape)
_logger.info(f"Feature maps shape: {output_shapes}")
_logger.info(f"Input shape : {input_shape}")
def gradient_function(
receptive_field_masks: List[torch.Tensor], ) -> List[np.ndarray]:
grads = []
for fm, rf_mask in enumerate(receptive_field_masks):
input_tensor = torch.zeros(*shape)
# input_tensor = torch.autograd.Variable(
# input_tensor, requires_grad=True
# )
input_tensor.requires_grad_(True)
model.zero_grad()
_ = model(input_tensor)
fm = torch.mean(model.feature_maps[fm], 1, keepdim=True)
fake_loss = fm * rf_mask
fake_loss = torch.mean(fake_loss)
fake_loss.backward()
grads.append(input_tensor.grad.detach().numpy())
return grads
return gradient_function, input_shape, output_shapes
def _prepare_gradient_func(
self, input_shape: ImageShape, input_tensor: str, output_tensors: List[str]
) -> Tuple[Callable, GridShape, List[GridShape]]:
"""
Computes gradient function and additional parameters. Note
that the receptive field parameters like stride or size, do not
depend on input image shape. However, if the RF of original network
is bigger than input_shape this method will fail. Hence it is
recommended to increase the input shape.
:param input_shape: shape of the input image. Used in @model_func.
:param input_tensor: name of the input image tensor.
:param output_tensors: a list of names of the target
feature map tensors.
:returns
gradient_function: a function which returns gradient w.r.t. to
the input image.
input_shape: a shape of the input image tensor.
output_shape: a list shapes of the output feature map tensors.
"""
if self._session is not None:
tf.reset_default_graph()
self._session.close()
with tf.Graph().as_default() as graph:
with tf.variable_scope("", reuse=tf.AUTO_REUSE):
# this function will create default graph
_ = self._model_func(ImageShape(*input_shape))
# default_graph = tf.get_default_graph()
# get graph tensors by names
input_tensor = graph.get_operation_by_name(input_tensor).outputs[0]
input_shape = _get_tensor_shape(input_tensor)
grads = []
receptive_field_masks = []
output_shapes = []
for output_tensor in output_tensors:
output_tensor = graph.get_operation_by_name(output_tensor).outputs[0]
# shapes
output_shape = _get_tensor_shape(output_tensor)
output_shape = (1, output_shape[1], output_shape[2], 1)
output_shapes.append(output_shape)
# define loss function
receptive_field_mask = tf.placeholder(
tf.float32, shape=output_shape, name="grid"
)
grad = _define_fm_gradient(
input_tensor, output_tensor, receptive_field_mask
)
grads.append(grad)
receptive_field_masks.append(receptive_field_mask)
_logger.info(f"Feature maps shape: {output_shapes}")
_logger.info(f"Input shape : {input_shape}")
self._session = tf.Session(graph=graph)
self._session.run(tf.global_variables_initializer())
def gradient_fn(fm_masks, input_image):
fetch_dict = {
mask_t: mask_np
for mask_t, mask_np in zip(receptive_field_masks, fm_masks)
}
fetch_dict[input_tensor] = input_image
return self._session.run(grads, feed_dict=fetch_dict)
return (
gradient_fn,
GridShape(*input_shape),
[GridShape(*output_shape) for output_shape in output_shapes],
)