Initial media depth project backup
This commit is contained in:
@@ -0,0 +1,121 @@
|
||||
# MIT License
|
||||
|
||||
# Copyright (c) 2022 Intelligent Systems Lab Org
|
||||
|
||||
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
# of this software and associated documentation files (the "Software"), to deal
|
||||
# in the Software without restriction, including without limitation the rights
|
||||
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
# copies of the Software, and to permit persons to whom the Software is
|
||||
# furnished to do so, subject to the following conditions:
|
||||
|
||||
# The above copyright notice and this permission notice shall be included in all
|
||||
# copies or substantial portions of the Software.
|
||||
|
||||
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
# SOFTWARE.
|
||||
|
||||
# File author: Shariq Farooq Bhat
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def log_binom(n, k, eps=1e-7):
|
||||
""" log(nCk) using stirling approximation """
|
||||
n = n + eps
|
||||
k = k + eps
|
||||
return n * torch.log(n) - k * torch.log(k) - (n-k) * torch.log(n-k+eps)
|
||||
|
||||
|
||||
class LogBinomial(nn.Module):
|
||||
def __init__(self, n_classes=256, act=torch.softmax):
|
||||
"""Compute log binomial distribution for n_classes
|
||||
|
||||
Args:
|
||||
n_classes (int, optional): number of output classes. Defaults to 256.
|
||||
"""
|
||||
super().__init__()
|
||||
self.K = n_classes
|
||||
self.act = act
|
||||
self.register_buffer('k_idx', torch.arange(
|
||||
0, n_classes).view(1, -1, 1, 1))
|
||||
self.register_buffer('K_minus_1', torch.Tensor(
|
||||
[self.K-1]).view(1, -1, 1, 1))
|
||||
|
||||
def forward(self, x, t=1., eps=1e-4):
|
||||
"""Compute log binomial distribution for x
|
||||
|
||||
Args:
|
||||
x (torch.Tensor - NCHW): probabilities
|
||||
t (float, torch.Tensor - NCHW, optional): Temperature of distribution. Defaults to 1..
|
||||
eps (float, optional): Small number for numerical stability. Defaults to 1e-4.
|
||||
|
||||
Returns:
|
||||
torch.Tensor -NCHW: log binomial distribution logbinomial(p;t)
|
||||
"""
|
||||
if x.ndim == 3:
|
||||
x = x.unsqueeze(1) # make it nchw
|
||||
|
||||
one_minus_x = torch.clamp(1 - x, eps, 1)
|
||||
x = torch.clamp(x, eps, 1)
|
||||
y = log_binom(self.K_minus_1, self.k_idx) + self.k_idx * \
|
||||
torch.log(x) + (self.K - 1 - self.k_idx) * torch.log(one_minus_x)
|
||||
return self.act(y/t, dim=1)
|
||||
|
||||
|
||||
class ConditionalLogBinomial(nn.Module):
|
||||
def __init__(self, in_features, condition_dim, n_classes=256, bottleneck_factor=2, p_eps=1e-4, max_temp=50, min_temp=1e-7, act=torch.softmax):
|
||||
"""Conditional Log Binomial distribution
|
||||
|
||||
Args:
|
||||
in_features (int): number of input channels in main feature
|
||||
condition_dim (int): number of input channels in condition feature
|
||||
n_classes (int, optional): Number of classes. Defaults to 256.
|
||||
bottleneck_factor (int, optional): Hidden dim factor. Defaults to 2.
|
||||
p_eps (float, optional): small eps value. Defaults to 1e-4.
|
||||
max_temp (float, optional): Maximum temperature of output distribution. Defaults to 50.
|
||||
min_temp (float, optional): Minimum temperature of output distribution. Defaults to 1e-7.
|
||||
"""
|
||||
super().__init__()
|
||||
self.p_eps = p_eps
|
||||
self.max_temp = max_temp
|
||||
self.min_temp = min_temp
|
||||
self.log_binomial_transform = LogBinomial(n_classes, act=act)
|
||||
bottleneck = (in_features + condition_dim) // bottleneck_factor
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Conv2d(in_features + condition_dim, bottleneck,
|
||||
kernel_size=1, stride=1, padding=0),
|
||||
nn.GELU(),
|
||||
# 2 for p linear norm, 2 for t linear norm
|
||||
nn.Conv2d(bottleneck, 2+2, kernel_size=1, stride=1, padding=0),
|
||||
nn.Softplus()
|
||||
)
|
||||
|
||||
def forward(self, x, cond):
|
||||
"""Forward pass
|
||||
|
||||
Args:
|
||||
x (torch.Tensor - NCHW): Main feature
|
||||
cond (torch.Tensor - NCHW): condition feature
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output log binomial distribution
|
||||
"""
|
||||
pt = self.mlp(torch.concat((x, cond), dim=1))
|
||||
p, t = pt[:, :2, ...], pt[:, 2:, ...]
|
||||
|
||||
p = p + self.p_eps
|
||||
p = p[:, 0, ...] / (p[:, 0, ...] + p[:, 1, ...])
|
||||
|
||||
t = t + self.p_eps
|
||||
t = t[:, 0, ...] / (t[:, 0, ...] + t[:, 1, ...])
|
||||
t = t.unsqueeze(1)
|
||||
t = (self.max_temp - self.min_temp) * t + self.min_temp
|
||||
|
||||
return self.log_binomial_transform(p, t)
|
||||
Reference in New Issue
Block a user