# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [default]
#     language: python
#     name: python3
# ---

# ## Gaussian Network Moments (GNMs)
#
# Let $\mathbf{x}\sim\mathcal{N}\left(\mathbf{\mu}, \mathbf{\Sigma}\right)$ and $q(x) = \max(x, 0)$ where $\Phi(x)$ and $\varphi(x)$ are the CDF and PDF of the normal distribution,
#
# $\mathbb{E}\left[q(\mathbf{x})\right] = \mathbf{\mu}\odot\Phi\left(\mathbf{\mu}\oslash\mathbf{\sigma}\right) + \mathbf{\sigma}\odot\varphi\left(\mathbf{\mu}\oslash\mathbf{\sigma}\right)$
# where $\mathbf{\sigma} = \sqrt{\text{diag}\left(\mathbf{\Sigma}\right)}$ with
# $\odot$ and $\oslash$ as element-wise product and division.
#
# $\mathbb{E}\left[q^2(\mathbf{x})\right] = 
# \left(\mathbf{\mu}^2+\mathbf{\sigma}^2\right) \odot \Phi\left(\mathbf{\mu}\oslash\mathbf{\sigma}\right) + \mathbf{\mu} \odot \mathbf{\sigma} \odot \varphi\left(\mathbf{\mu}\oslash\mathbf{\sigma}\right)$
# where $\text{var}\left[q(\mathbf{x})\right] = \mathbb{E}\left[q^2(\mathbf{x})\right] - \mathbb{E}\left[q(\mathbf{x})\right]^2$
#
# $\left.\mathbb{E}\left[q(\mathbf{x})q(\mathbf{x})^\top\right]\right|_{\mathbf{\mu} = \mathbf{0}} = c\left(\mathbf{\Sigma}\oslash\mathbf{\sigma}\mathbf{\sigma}^\top\right) \odot \mathbf{\sigma}\mathbf{\sigma}^\top$
# where $c(x) = \frac{1}{2\pi}\left(x\cos^{-1}(-x)+\sqrt{1-x^2}\right)$
# (Note: $\left|c(x) - \Phi(x - 1)\right| < 0.0241$)
#
# $\text{cov}\left[q(\mathbf{x})\right] = \mathbb{E}\left[q(\mathbf{x})q(\mathbf{x})^\top\right] - \mathbb{E}\left[q(\mathbf{x})\right]\mathbb{E}\left[q(\mathbf{x})\right]^\top$
# where $\left.\text{cov}\left[q(\mathbf{x})\right]\right|_{\mathbf{\mu} = \mathbf{0}} = \left.\mathbb{E}\left[q(\mathbf{x})q(\mathbf{x})^\top\right]\right|_{\mathbf{\mu} = \mathbf{0}} - \frac{1}{2\pi}\mathbf{\sigma}\mathbf{\sigma}^\top$

# +
import sys
import torch
import unittest
from torch import nn
from pprint import pprint
from types import ModuleType
import network_moments.torch as nm

seed = 77  # for reproducability

def traverse(obj, exclude=[]):
    data = []
    if type(obj) is not ModuleType:
        return data
    for e in dir(obj):
        if not e.startswith('_') and all(e != s for s in exclude):
            sub = traverse(obj.__dict__[e], exclude)
            data.append(e if len(sub) == 0 else {e:sub})
    return data

print(nm.__doc__)
print('Network Moments Structure:')
pprint(traverse(nm.gaussian, exclude=['tests', 'general']))
# -

# ### Testing the tightness of the expressions on using tests

runner = unittest.TextTestRunner(sys.stdout, verbosity=2)
load = unittest.TestLoader().loadTestsFromModule
result = runner.run(unittest.TestSuite([
    load(nm.gaussian.affine.tests),
    load(nm.gaussian.relu.tests),
]))

# ### Testing the tightness of the expressions on affine-ReLU-affine networks

rand = nm.utils.rand
gnm = nm.gaussian.affine_relu_affine
print(gnm.special_variance.__doc__)

# +
length = 3
count = 1000000
dtype = torch.float64
device = torch.device('cpu', 0)
torch.manual_seed(seed)

# input mean and covariance
mu = torch.randn(length, dtype=dtype, device=device)
cov = rand.definite(length, dtype=dtype, device=device,
                    positive=True, semi=False, norm=1.0)

# variables
A = torch.randn(length, length, dtype=dtype, device=device)
c1 = -A.matmul(mu)  # torch.randn(length, dtype=dtype)
B = torch.randn(length, length, dtype=dtype, device=device)
c2 = torch.randn(length, dtype=dtype, device=device)

# analytical output mean and variance
out_mu = gnm.mean(mu, cov, A, c1, B, c2)
out_var = gnm.special_variance(cov, A, B)

# Monte-Carlo estimation of the output mean and variance
normal = torch.distributions.MultivariateNormal(mu, cov)
samples = normal.sample((count,))
out_samples = samples.matmul(A.t()) + c1
out_samples = torch.max(out_samples, torch.zeros([], dtype=dtype, device=device))
out_samples = out_samples.matmul(B.t()) + c2
mc_mu = torch.mean(out_samples, dim=0)
mc_var = torch.var(out_samples, dim=0)

# printing the ratios
print('Monte-Carlo mean / Analytical mean:')
print((mc_mu / out_mu).cpu().numpy())
print('Monte-Carlo variance / Analytical variance:')
print((mc_var / out_var).cpu().numpy())
# -

# ### Linearization

# +
batch = 1
num_classes = 10
image_size = (28, 28)
dtype = torch.float64
device = torch.device('cpu', 0)

size = torch.prod(torch.tensor(image_size)).item()
x = torch.rand(batch, *image_size, dtype=dtype, device=device)
model = nn.Sequential(
    nm.utils.flatten,
    nn.Linear(size, num_classes),
)
model.type(dtype)
if device.type != 'cpu':
    model.cuda(device.index)

jac, bias = nm.utils.linearize(model, x)
A = list(model.children())[1].weight
print('Tightness of A (best is zero): {}'.format(
    torch.max(torch.abs(jac - A)).item()))

b = list(model.children())[1].bias
print('Tightness of b (best is zero): {}'.format(
    torch.max(torch.abs(bias - b)).item()))
# -

# ### Two-stage linearization

# +
count = 10000
num_classes = 10
image_size = (28, 28)
dtype = torch.float64
device = torch.device('cpu', 0)
gnm = nm.gaussian.affine_relu_affine

size = torch.prod(torch.tensor(image_size)).item()
x = torch.rand(1, *image_size, dtype=dtype, device=device)

# deep model
first_part = nn.Sequential(
    nm.utils.flatten,
    nn.Linear(size, 500),
    nn.ReLU(),
    nn.Linear(500, 500),
    nn.ReLU(),
    nn.Linear(500, 300),
)
first_part.type(dtype)
relu = nn.Sequential(
    nn.ReLU(),
)
relu.type(dtype)
second_part = nn.Sequential(
    nn.Linear(300, 100),
    nn.ReLU(),
    nn.Linear(100, num_classes),
)
second_part.type(dtype)
if device.type != 'cpu':
    first_part.cuda(device.index)
    relu.cuda(device.index)
    second_part.cuda(device.index)
def model(x):
    return second_part(relu(first_part(x)))

# variables
A, c1 = nm.utils.linearize(first_part, x)
B, c2 = nm.utils.linearize(second_part, relu(first_part(x)).detach())
x.requires_grad_(False)
A.squeeze_()
c1.squeeze_()
B.squeeze_()
c2.squeeze_()

# analytical output mean and variance
mean = x.view(-1)
covariance = rand.definite(size, norm=0.1, dtype=dtype, device=device)
out_mu = gnm.mean(mean, covariance, A, c1, B, c2)
out_var = gnm.special_variance(covariance, A, B)

# Monte-Carlo estimation of the output mean and variance
normal = torch.distributions.MultivariateNormal(mean, covariance)
samples = normal.sample((count,))
out_samples = model(samples.view(-1, *image_size)).detach()
mc_mu = torch.mean(out_samples, dim=0)
mc_var = torch.var(out_samples, dim=0)

# printing the ratios
print('Monte-Carlo mean / Analytical mean:')
print((mc_mu / out_mu).cpu().numpy())
print('Monte-Carlo variance / Analytical variance:')
print((mc_var / out_var).cpu().numpy())

StarcoderdataJupyterScriptsDedupFiltered