Source code for jactorch.nn.probability

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : probability.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 02/04/2018
#
# This file is part of Jacinle.
# Distributed under terms of the MIT license.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter

from jactorch.functional.probability import normalize_prob

__all__ = ['ProbabilityLinear', 'ProbabilityBilinear', 'ProbabilityNLinear']

"""
ProbabilityLinear layers are "self-normalized" linear transformations of probability distributions.
More specifically:
    - When the input is normalized along the last dimension (i.e., values sum up to 1), the output is also normalized.
    - Thus, the transformation can be interpreted as a probability mass transformation.
    - The linear transformation has only the weight matrix. The weights are parameterized as a log-softmax distribution.
"""


[docs] class ProbabilityLinear(nn.Linear):
[docs] def __init__(self, in_features, out_features, bias=False, norm=True): assert bias is False, 'Bias regularization for SOFTMAX is not implemented.' super().__init__(in_features, out_features, bias) self.norm = norm
[docs] def forward(self, input): weight = self._regulize_parameter(self.weight) output = F.linear(input, weight, None) if self.norm: output = normalize_prob(output) return output
def _regulize_parameter(self, p): return F.softmax(p, dim=0)
[docs] class ProbabilityBilinear(nn.Bilinear):
[docs] def __init__(self, in1_features, in2_features, out_features, bias=False, norm=True): assert bias is False, 'Bias regularization for SOFTMAX is not implemented.' super().__init__(in1_features, in2_features, out_features, bias) self.norm = norm
[docs] def forward(self, input1, input2): weight = self._regulize_parameter(self.weight) output = F.bilinear(input1, input2, weight, None) if self.norm: output = normalize_prob(output) return output
def _regulize_parameter(self, p): return F.softmax(p, dim=0)
[docs] class ProbabilityNLinear(nn.Module):
[docs] def __new__(cls, *nr_categories, bias=False, norm=True): if len(nr_categories) == 2: return ProbabilityLinear(*nr_categories, bias=bias, norm=norm) elif len(nr_categories) == 3: return ProbabilityBilinear(*nr_categories, bias=bias, norm=norm) else: return super().__new__(cls)
[docs] def __init__(self, *nr_categories, bias=False, norm=True): super().__init__() assert bias is False, 'Bias regularization for SOFTMAX is not implemented.' self.nr_categories = nr_categories self.weight = Parameter(torch.Tensor(nr_categories[-1], *nr_categories[:-1])) self.reset_parameters() self.norm = norm
[docs] def reset_parameters(self): stdv = 1. / math.sqrt(self.weight.size(1)) self.weight.data.uniform_(-stdv, stdv)
[docs] def forward(self, *inputs): f = self._regulize_parameter(self.weight) for i in reversed(inputs): f = (f * i).sum(dim=-1) if self.norm: f = normalize_prob(f) return f
def _regulize_parameter(self, p): return F.softmax(p, dim=0)