#! /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)