coord_map.py

"""
Determine spatial relationships between layers to relate their coordinates.
Coordinates are mapped from input-to-output (forward), but can
be mapped output-to-input (backward) by the inverse mapping too.
This helps crop and align feature maps among other uses.
"""

from __future__ import division
import numpy as np
from caffe import layers as L

PASS_THROUGH_LAYERS = ['AbsVal', 'BatchNorm', 'Bias', 'BNLL', 'Dropout',
                       'Eltwise', 'ELU', 'Log', 'LRN', 'Exp', 'MVN', 'Power',
                       'ReLU', 'PReLU', 'Scale', 'Sigmoid', 'Split', 'TanH',
                       'Threshold']


def conv_params(fn):
    """
    Extract the spatial parameters that determine the coordinate mapping:
    kernel size, stride, padding, and dilation.

    Implementation detail: Convolution, Deconvolution, and Im2col layers
    define these in the convolution_param message, while Pooling has its
    own fields in pooling_param. This method deals with these details to
    extract canonical parameters.
    """
    params = fn.params.get('convolution_param', fn.params)
    axis = params.get('axis', 1)
    ks = np.array(params['kernel_size'], ndmin=1)
    dilation = np.array(params.get('dilation', 1), ndmin=1)
    assert len({'pad_h', 'pad_w', 'kernel_h', 'kernel_w', 'stride_h',
                'stride_w'} & set(fn.params)) == 0, \
        'cropping does not support legacy _h/_w params'
    return (axis, np.array(params.get('stride', 1), ndmin=1),
            (ks - 1) * dilation + 1,
            np.array(params.get('pad', 0), ndmin=1))


def crop_params(fn):
    """
    Extract the crop layer parameters with defaults.
    """
    params = fn.params.get('crop_param', fn.params)
    axis = params.get('axis', 2)  # default to spatial crop for N, C, H, W
    offset = np.array(params.get('offset', 0), ndmin=1)
    return (axis, offset)


class UndefinedMapException(Exception):
    """
    Exception raised for layers that do not have a defined coordinate mapping.
    """
    pass


def coord_map(fn):
    """
    Define the coordinate mapping by its
    - axis
    - scale: output coord[i * scale] <- input_coord[i]
    - shift: output coord[i] <- output_coord[i + shift]
    s.t. the identity mapping, as for pointwise layers like ReLu, is defined by
    (None, 1, 0) since it is independent of axis and does not transform coords.
    """
    if fn.type_name in ['Convolution', 'Pooling', 'Im2col']:
        axis, stride, ks, pad = conv_params(fn)
        return axis, 1 / stride, (pad - (ks - 1) / 2) / stride
    elif fn.type_name == 'Deconvolution':
        axis, stride, ks, pad = conv_params(fn)
        return axis, stride, (ks - 1) / 2 - pad
    elif fn.type_name in PASS_THROUGH_LAYERS:
        return None, 1, 0
    elif fn.type_name == 'Crop':
        axis, offset = crop_params(fn)
        axis -= 1  # -1 for last non-coordinate dim.
        return axis, 1, - offset
    else:
        raise UndefinedMapException


class AxisMismatchException(Exception):
    """
    Exception raised for mappings with incompatible axes.
    """
    pass


def compose(base_map, next_map):
    """
    Compose a base coord map with scale a1, shift b1 with a further coord map
    with scale a2, shift b2. The scales multiply and the further shift, b2,
    is scaled by base coord scale a1.
    """
    ax1, a1, b1 = base_map
    ax2, a2, b2 = next_map
    if ax1 is None:
        ax = ax2
    elif ax2 is None or ax1 == ax2:
        ax = ax1
    else:
        raise AxisMismatchException
    return ax, a1 * a2, a1 * b2 + b1


def inverse(coord_map):
    """
    Invert a coord map by de-scaling and un-shifting;
    this gives the backward mapping for the gradient.
    """
    ax, a, b = coord_map
    return ax, 1 / a, -b / a


def coord_map_from_to(top_from, top_to):
    """
    Determine the coordinate mapping betweeen a top (from) and a top (to).
    Walk the graph to find a common ancestor while composing the coord maps for
    from and to until they meet. As a last step the from map is inverted.
    """
    # We need to find a common ancestor of top_from and top_to.
    # We'll assume that all ancestors are equivalent here (otherwise the graph
    # is an inconsistent state (which we could improve this to check for)).
    # For now use a brute-force algorithm.

    def collect_bottoms(top):
        """
        Collect the bottoms to walk for the coordinate mapping.
        The general rule is that all the bottoms of a layer can be mapped, as
        most layers have the same coordinate mapping for each bottom.
        Crop layer is a notable exception. Only the first/cropped bottom is
        mappable; the second/dimensions bottom is excluded from the walk.
        """
        bottoms = top.fn.inputs
        if top.fn.type_name == 'Crop':
            bottoms = bottoms[:1]
        return bottoms

    # walk back from top_from, keeping the coord map as we go
    from_maps = {top_from: (None, 1, 0)}
    frontier = {top_from}
    while frontier:
        top = frontier.pop()
        try:
            bottoms = collect_bottoms(top)
            for bottom in bottoms:
                from_maps[bottom] = compose(from_maps[top], coord_map(top.fn))
                frontier.add(bottom)
        except UndefinedMapException:
            pass

    # now walk back from top_to until we hit a common blob
    to_maps = {top_to: (None, 1, 0)}
    frontier = {top_to}
    while frontier:
        top = frontier.pop()
        if top in from_maps:
            return compose(to_maps[top], inverse(from_maps[top]))
        try:
            bottoms = collect_bottoms(top)
            for bottom in bottoms:
                to_maps[bottom] = compose(to_maps[top], coord_map(top.fn))
                frontier.add(bottom)
        except UndefinedMapException:
            continue

    # if we got here, we did not find a blob in common
    raise RuntimeError('Could not compute map between tops; are they '
                       'connected by spatial layers?')


def crop(top_from, top_to):
    """
    Define a Crop layer to crop a top (from) to another top (to) by
    determining the coordinate mapping between the two and net spec'ing
    the axis and shift parameters of the crop.
    """
    ax, a, b = coord_map_from_to(top_from, top_to)
    assert (a == 1).all(), 'scale mismatch on crop (a = {})'.format(a)
    assert (b <= 0).all(), 'cannot crop negative offset (b = {})'.format(b)
    assert (np.round(b) == b).all(), 'cannot crop noninteger offset ' \
        '(b = {})'.format(b)
    return L.Crop(top_from, top_to,
                  crop_param=dict(axis=ax + 1,  # +1 for first cropping dim.
                                  offset=list(-np.round(b).astype(int))))