# LICENSE: EXTERNAL
# MIT License
#
# Copyright (c) 2021 Chin-Wei Huang
#
# 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.
from typing import Optional, Tuple, List, Union
import torch
import torch.nn.functional as F
from torch.distributions.normal import Normal
from ._cpflows import SequentialFlow, DeepConvexFlow
[docs]class DeepConvexNet(DeepConvexFlow):
r"""
Class that takes a partially input convex neural network (picnn) as input
and equips it with functions of logdet computation (both estimation and
exact computation)
This class is based on DeepConvexFlow of the CP-Flow
repo (https://github.com/CW-Huang/CP-Flow)
For details of the logdet estimator, see
``Convex potential flows: Universal probability distributions
with optimal transport and convex optimization``
Parameters
----------
picnn
A partially input convex neural network (picnn)
dim
Dimension of the input
is_energy_score
Indicates if energy score is used as the objective function
If yes, the network is not required to be strictly convex,
so we can just use the picnn
otherwise, a quadratic term is added to the output of picnn
to render it strictly convex
m1
Dimension of the Krylov subspace of the Lanczos tridiagonalization
used in approximating H of logdet(H)
m2
Iteration number of the conjugate gradient algorithm
used to approximate logdet(H)
rtol
relative tolerance of the conjugate gradient algorithm
atol
absolute tolerance of the conjugate gradient algorithm
"""
def __init__(
self,
picnn: torch.nn.Module,
dim: int,
is_energy_score: bool = False,
estimate_logdet: bool = False,
m1: int = 10,
m2: Optional[int] = None,
rtol: float = 0.0,
atol: float = 1e-3,
) -> None:
super().__init__(
picnn,
dim,
m1=m1,
m2=m2,
rtol=rtol,
atol=atol,
)
self.picnn = self.icnn
self.is_energy_score = is_energy_score
self.estimate_logdet = estimate_logdet
[docs] def get_potential(
self, x: torch.Tensor, context: Optional[torch.Tensor] = None
) -> torch.Tensor:
n = x.size(0)
output = self.picnn(x, context)
if self.is_energy_score:
return output
else:
return (
F.softplus(self.w1) * output
+ F.softplus(self.w0)
* (x.view(n, -1) ** 2).sum(1, keepdim=True)
/ 2
)
[docs]class SequentialNet(SequentialFlow):
r"""
Class that combines a list of DeepConvexNet and ActNorm layers and provides
energy score computation.
This class is based on SequentialFlow of the CP-Flow repo
(https://github.com/CW-Huang/CP-Flow)
Parameters
----------
networks
list of DeepConvexNet and/or ActNorm instances
"""
def __init__(self, networks: List[torch.nn.Module]) -> None:
super().__init__(networks)
self.networks = self.flows
[docs] def forward(
self, x: torch.Tensor, context: Optional[torch.Tensor] = None
) -> torch.Tensor:
for network in self.networks:
if isinstance(network, DeepConvexNet):
x = network.forward(x, context=context)
else:
x = network.forward(x)
return x
[docs] def es_sample(
self, hidden_state: torch.Tensor, dimension: int
) -> torch.Tensor:
"""
Auxiliary function for energy score computation Drawing samples
conditioned on the hidden state.
Parameters
----------
hidden_state
hidden_state which the samples conditioned
on (num_samples, hidden_size)
dimension
dimension of the input
Returns
-------
samples
samples drawn (num_samples, dimension)
"""
num_samples = hidden_state.shape[0]
zero = torch.tensor(
0, dtype=hidden_state.dtype, device=hidden_state.device
)
one = torch.ones_like(zero)
standard_normal = Normal(zero, one)
samples = self.forward(
standard_normal.sample([num_samples * dimension]).view(
num_samples, dimension
),
context=hidden_state,
)
return samples
[docs] def energy_score(
self,
z: torch.Tensor,
hidden_state: torch.Tensor,
es_num_samples: int = 50,
beta: float = 1.0,
) -> torch.Tensor:
"""
Computes the (approximated) energy score sum_i ES(g,z_i), where
ES(g,z_i) =
-1/(2*es_num_samples^2) * sum_{w,w'} ||w-w'||_2^beta
+ 1/es_num_samples * sum_{w''} ||w''-z_i||_2^beta,
w's are samples drawn from the
quantile function g(., h_i) (gradient of picnn),
h_i is the hidden state associated with z_i,
and es_num_samples is the number of samples drawn
for each of w, w', w'' in energy score approximation
Parameters
----------
z
Observations (numel_batch, dimension)
hidden_state
Hidden state (numel_batch, hidden_size)
es_num_samples
Number of samples drawn for each of w, w', w''
in energy score approximation
beta
Hyperparameter of the energy score, see the formula above
Returns
-------
loss
energy score (numel_batch)
"""
numel_batch, dimension = z.shape[0], z.shape[1]
hidden_state_repeat = hidden_state.repeat_interleave(
repeats=es_num_samples, dim=0
)
w = self.es_sample(hidden_state_repeat, dimension)
w_prime = self.es_sample(hidden_state_repeat, dimension)
first_term = (
torch.norm(
w.view(numel_batch, 1, es_num_samples, dimension)
- w_prime.view(numel_batch, es_num_samples, 1, dimension),
dim=-1,
)
** beta
)
mean_first_term = torch.mean(first_term.view(numel_batch, -1), dim=-1)
# since both tensors are huge (numel_batch*es_num_samples, dimension),
# delete to free up GPU memories
del w, w_prime
z_repeat = z.repeat_interleave(repeats=es_num_samples, dim=0)
w_bar = self.es_sample(hidden_state_repeat, dimension)
second_term = (
torch.norm(
w_bar.view(numel_batch, es_num_samples, dimension)
- z_repeat.view(numel_batch, es_num_samples, dimension),
dim=-1,
)
** beta
)
mean_second_term = torch.mean(
second_term.view(numel_batch, -1), dim=-1
)
loss = -0.5 * mean_first_term + mean_second_term
return loss