prob_average.py

"""
Written by Hokuto Kagaya, Dec. 2015

ver 0.1

"""

import numpy

from chainer import cuda
from chainer import function
from chainer.utils import type_check


class ProbAverage(function.Function):

    """Average probabilities."""

    def __init__(self, ctof=None, ftoc=None, fine_size=None):
        self.ctof = ctof
        self.ftoc = ftoc
        self.fine_size = fine_size

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() == 1 + len(self.ctof))

        c_type = in_types[0]
        fs_type = in_types[1:]

        # type_check.expect(fs_type.size() > 0)

    def forward(self, inputs):
        xp = cuda.get_array_module(*inputs)
        xc = inputs[0]
        xhs = inputs[1:]

        batch_size = xc.shape[0]
        coarse_size = xc[0].shape[0]

        y = xp.zeros((batch_size, self.fine_size), dtype=numpy.float32)
        for b in xrange(batch_size):
            xcb = xc[b]
            for c in xrange(coarse_size):
                each_fine_size = len(self.ctof[c])
                for k in xrange(each_fine_size):
                    y[b][self.ctof[c][k]] += xhs[c][b][k] * xcb[c]

        return y,


    def backward(self, xs, gy):
        xp = cuda.get_array_module(*xs)

        _gy = gy[0]
        assert _gy.shape[1] == self.fine_size

        ret_to_backward = [xp.zeros((_gy.shape[0], len(self.ctof)), dtype=numpy.float32) if i == 0 else xp.zeros((_gy.shape[0], len(self.ctof[i-1])), dtype=numpy.float32) for i in xrange(len(self.ctof)+1)]
        # print 

        for i in xrange(_gy.shape[0]):
            for f in xrange(self.fine_size):
                f_ftoc = self.ftoc[f]
                for c in f_ftoc:
                    c_ctof = self.ctof[c]
                    ret_to_backward[c+1][i][c_ctof.index(f)] += _gy[i][f] * xs[0][i][c]
                    ret_to_backward[0][i][c] += _gy[i][f] * xs[c+1][i][c_ctof.index(f)]

        return tuple(ret_to_backward)

def prob_average(xc, xfs, ctof, ftoc, fine_size):
    """Average probabilities.

    Args:
        xs (tuple of Variables): Variables to be concatenated.
        axis (int): Axis that the input arrays are concatenated along.

    Returns:
        ~chainer.Variable: Output variable.

    """
    return ProbAverage(ctof, ftoc, fine_size)(xc, *xfs)