wassname/attentive-neural-processes
1
0
Fork 0
mirror of https://github.com/wassname/attentive-neural-processes.git synced 2026-07-20 12:10:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

small fixes, working

Browse Source
This commit is contained in:
wassname
2019-11-02 11:50:39 +08:00
parent ce08fd9bb3
commit 0dec1eb3bc
9 changed files with 645 additions and 3017 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -1,3 +1,5 @@
data/
# Created by https://www.gitignore.io/api/code,linux,macos,python,windows,jupyternotebook,jupyternotebooks
# Edit at https://www.gitignore.io/?templates=code,linux,macos,python,windows,jupyternotebook,jupyternotebooks
anp_1d_regression.ipynb
+105 -3017
View File
File diff suppressed because one or more lines are too long
readme.md
+1
View File
@@ -0,0 +1 @@
- Download data from https://www.kaggle.com/jeanmidev/smart-meters-in-london/version/11
src/__init__.py
View File
src/data/__init__.py
View File
src/data/gp_curves.py
+195
View File
@@ -0,0 +1,195 @@
import pandas as pd
import numpy as np
import collections
import torch
# The (A)NP takes as input a `NPRegressionDescription` namedtuple with fields:
# `query`: a tuple containing ((context_x, context_y), target_x)
# `target_y`: a tensor containing the ground truth for the targets to be
# predicted
# `num_total_points`: A vector containing a scalar that describes the total
# number of datapoints used (context + target)
# `num_context_points`: A vector containing a scalar that describes the number
# of datapoints used as context
# The GPCurvesReader returns the newly sampled data in this format at each
# iteration
NPRegressionDescription = collections.namedtuple(
"NPRegressionDescription",
("query", "target_y", "num_total_points", "num_context_points"),
)
class GPCurvesReader(object):
"""Generates curves using a Gaussian Process (GP).
Supports vector inputs (x) and vector outputs (y). Kernel is
mean-squared exponential, using the x-value l2 coordinate distance scaled by
some factor chosen randomly in a range. Outputs are independent gaussian
processes.
"""
def __init__(
self,
batch_size,
max_num_context,
x_size=1,
y_size=1,
l1_scale=0.6,
sigma_scale=1.0,
random_kernel_parameters=True,
testing=False,
):
"""Creates a regression dataset of functions sampled from a GP.
Args:
batch_size: An integer.
max_num_context: The max number of observations in the context.
x_size: Integer >= 1 for length of "x values" vector.
y_size: Integer >= 1 for length of "y values" vector.
l1_scale: Float; typical scale for kernel distance function.
sigma_scale: Float; typical scale for variance.
random_kernel_parameters: If `True`, the kernel parameters (l1 and sigma)
will be sampled uniformly within [0.1, l1_scale] and [0.1, sigma_scale].
testing: Boolean that indicates whether we are testing. If so there are
more targets for visualization.
"""
self._batch_size = batch_size
self._max_num_context = max_num_context
self._x_size = x_size
self._y_size = y_size
self._l1_scale = l1_scale
self._sigma_scale = sigma_scale
self._random_kernel_parameters = random_kernel_parameters
self._testing = testing
def _gaussian_kernel(self, xdata, l1, sigma_f, sigma_noise=2e-2):
"""Applies the Gaussian kernel to generate curve data.
Args:
xdata: Tensor of shape [B, num_total_points, x_size] with
the values of the x-axis data.
l1: Tensor of shape [B, y_size, x_size], the scale
parameter of the Gaussian kernel.
sigma_f: Tensor of shape [B, y_size], the magnitude
of the std.
sigma_noise: Float, std of the noise that we add for stability.
Returns:
The kernel, a float tensor of shape
[B, y_size, num_total_points, num_total_points].
"""
num_total_points = xdata.shape[1]
# Expand and take the difference
xdata1 = xdata.unsqueeze(1) # [B, 1, num_total_points, x_size]
xdata2 = xdata.unsqueeze(2) # [B, num_total_points, 1, x_size]
diff = xdata1 - xdata2 # [B, num_total_points, num_total_points, x_size]
# [B, y_size, num_total_points, num_total_points, x_size]
norm = (diff[:, None, :, :, :] / l1[:, :, None, None, :]) ** 2
norm = torch.sum(norm, -1) # [B, data_size, num_total_points, num_total_points]
# [B, y_size, num_total_points, num_total_points]
kernel = ((sigma_f) ** 2)[:, :, None, None] * torch.exp(-0.5 * norm)
# Add some noise to the diagonal to make the cholesky work.
kernel += (sigma_noise ** 2) * torch.eye(num_total_points)
return kernel
def generate_curves(self):
"""Builds the op delivering the data.
Generated functions are `float32` with x values between -2 and 2.
Returns:
A `CNPRegressionDescription` namedtuple.
"""
num_context = int(np.random.rand() * (self._max_num_context - 3) + 3)
# If we are testing we want to have more targets and have them evenly
# distributed in order to plot the function.
if self._testing:
num_target = 400
num_total_points = num_target
x_values = (
torch.arange(-2, 2, 1.0 / 100).unsqueeze(0).repeat(self._batch_size, 1)
)
x_values = x_values.unsqueeze(-1)
# During training the number of target points and their x-positions are
# selected at random
else:
num_target = int(np.random.rand() * (self._max_num_context - num_context))
num_total_points = num_context + num_target
x_values = (
torch.rand((self._batch_size, num_total_points, self._x_size)) * 4 - 2
)
# Set kernel parameters
# Either choose a set of random parameters for the mini-batch
if self._random_kernel_parameters:
l1 = (
torch.rand((self._batch_size, self._y_size, self._x_size))
* (self._l1_scale - 0.1)
+ 0.1
)
sigma_f = (
torch.rand((self._batch_size, self._y_size)) * (self._sigma_scale - 0.1)
+ 0.1
)
# Or use the same fixed parameters for all mini-batches
else:
l1 = (
torch.ones((self._batch_size, self._y_size, self._x_size))
* self._l1_scale
)
sigma_f = torch.ones((self._batch_size, self._y_size)) * self._sigma_scale
# Pass the x_values through the Gaussian kernel
# [batch_size, y_size, num_total_points, num_total_points]
kernel = self._gaussian_kernel(x_values, l1, sigma_f)
# Calculate Cholesky, using double precision for better stability:
cholesky = torch.cholesky(kernel)
# Sample a curve
# [batch_size, y_size, num_total_points, 1]
y_values = torch.matmul(
cholesky, torch.randn((self._batch_size, self._y_size, num_total_points, 1))
)
# [batch_size, num_total_points, y_size]
y_values = y_values.squeeze(3)
y_values = y_values.permute(0, 2, 1)
if self._testing:
# Select the targets
target_x = x_values
target_y = y_values
# Select the observations
idx = torch.randperm(num_target)
context_x = x_values[:, idx[:num_context]]
context_y = y_values[:, idx[:num_context]]
else:
# Select the targets which will consist of the context points as well as
# some new target points
target_x = x_values[:, : num_target + num_context, :]
target_y = y_values[:, : num_target + num_context, :]
# Select the observations
context_x = x_values[:, :num_context, :]
context_y = y_values[:, :num_context, :]
query = ((context_x, context_y), target_x)
return NPRegressionDescription(
query=query,
target_y=target_y,
num_total_points=target_x.shape[1],
num_context_points=num_context,
)
src/models/__init__.py
View File
src/models/model.py
+119
View File
@@ -0,0 +1,119 @@
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
import math
from .modules import LatentEncoder, DeterministicEncoder, Decoder
def log_prob_sigma(value, loc, log_scale):
"""A slightly more stable (not confirmed yet) log prob taking in log_var instead of scale.
modified from https://github.com/pytorch/pytorch/blob/2431eac7c011afe42d4c22b8b3f46dedae65e7c0/torch/distributions/normal.py#L65
"""
var = torch.exp(log_scale * 2)
return (
-((value - loc) ** 2) / (2 * var) - log_scale - math.log(math.sqrt(2 * math.pi))
)
def kl_loss_var(prior_mu, log_var_prior, post_mu, log_var_post):
"""
Analytical KLD for two gaussians, taking in log_variance instead of scale ( given variance=scale**2) for more stable gradients
For version using scale see https://github.com/pytorch/pytorch/blob/master/torch/distributions/kl.py#L398
"""
var_ratio_log = log_var_post - log_var_prior
kl_div = (
(var_ratio_log.exp() + (post_mu - prior_mu) ** 2) / log_var_prior.exp()
- 1.0
- var_ratio_log
)
kl_div = 0.5 * kl_div
return kl_div
class LatentModel(nn.Module):
def __init__(
self,
x_dim,
y_dim,
hidden_dim=32,
latent_dim=32,
latent_enc_self_attn_type="multihead",
det_enc_self_attn_type="multihead",
det_enc_cross_attn_type="multihead",
n_latent_encoder_layers=3,
n_det_encoder_layers=3,
n_decoder_layers=3,
num_heads=8,
dropout=0,
):
super(LatentModel, self).__init__()
self._latent_encoder = LatentEncoder(
x_dim + y_dim,
hidden_dim=hidden_dim,
latent_dim=latent_dim,
self_attention_type=latent_enc_self_attn_type,
n_encoder_layers=n_latent_encoder_layers,
dropout=dropout,
n_heads=num_heads,
)
self._deterministic_encoder = DeterministicEncoder(
x_dim + y_dim,
x_dim,
hidden_dim=hidden_dim,
self_attention_type=det_enc_self_attn_type,
cross_attention_type=det_enc_cross_attn_type,
n_d_encoder_layers=n_det_encoder_layers,
dropout=dropout,
n_heads=num_heads,
)
self._decoder = Decoder(
x_dim,
y_dim,
hidden_dim=hidden_dim,
n_decoder_layers=n_decoder_layers,
dropout=dropout,
)
def forward(self, context_x, context_y, target_x, target_y=None):
num_targets = target_x.size(1)
dist_prior, log_var_prior = self._latent_encoder(context_x, context_y)
if target_y is not None:
dist_post, log_var_post = self._latent_encoder(target_x, target_y)
z = dist_post.rsample()
else:
z = (
dist_prior.loc
) # instead of sampling, in test mode take the mean, this will make it more deterministic
z = z.unsqueeze(1).repeat(1, num_targets, 1) # [B, T_target, H]
r = self._deterministic_encoder(
context_x, context_y, target_x
) # [B, T_target, H]
dist, log_sigma = self._decoder(r, z, target_x)
if target_y is not None:
# Log likelihood has shape (batch_size, num_target, y_dim).
log_p = log_prob_sigma(target_y, dist.loc, log_sigma).mean(-1)
# KL has shape (batch_size, r_dim)
kl_loss = kl_loss_var(
dist_prior.loc, log_var_prior, dist_post.loc, log_var_post
).mean(-1)
kl_loss = kl_loss[:, None].expand(log_p.shape)
loss = (kl_loss - log_p).mean()
else:
log_p = None
kl_loss = None
loss = None
return dist.rsample(), kl_loss, loss, dist.scale
src/models/modules.py
+223
View File
@@ -0,0 +1,223 @@
import torch
from torch import nn
import torch.nn.functional as F
import math
class NPBlockRelu2d(nn.Module):
"""Block for Neural Processes."""
def __init__(self, in_channels, out_channels, dropout=0, norm=True):
super(NPBlockRelu2d, self).__init__()
self.linear = nn.Linear(in_channels, out_channels)
self.act = nn.ReLU()
self.dropout = nn.Dropout2d(dropout)
self.norm = nn.BatchNorm2d(out_channels) if norm else False
def forward(self, x):
# x.shape is (Batch, Sequence, Channels)
# We pass a linear over it which operates on the Channels
x = self.act(self.linear(x))
# Now we want to apply batchnorm and dropout to the channels. So we put it in shape
# (Batch, Channels, Sequence, None) so we can use Dropout2d
x = x.permute(0, 2, 1)[:, :, :, None]
if self.norm:
x = self.norm(x)
x = self.dropout(x)
return x[:, :, :, 0].permute(0, 2, 1)
def block_relu(in_dim, out_dim, dropout=0, inplace=False):
return nn.Sequential(
nn.Linear(in_dim, out_dim),
nn.ReLU(inplace=inplace),
nn.BatchNorm1d(out_dim),
nn.Dropout(dropout, inplace=inplace),
)
class Attention(nn.Module):
def __init__(self, hidden_dim, attention_type, n_heads=8, dropout=0):
super().__init__()
if attention_type == "uniform":
self._attention_func = self._uniform_attention
elif attention_type == "laplace":
self._attention_func = self._laplace_attention
elif attention_type == "dot":
self._attention_func = self._dot_attention
elif attention_type == "multihead":
self._mattn = torch.nn.MultiheadAttention(
hidden_dim, n_heads, bias=False, dropout=dropout
)
self._attention_func = self._pytorch_multihead_attention
self.n_heads = n_heads
else:
raise NotImplementedError
def forward(self, k, v, q):
rep = self._attention_func(k, v, q)
return rep
def _uniform_attention(self, k, v, q):
total_points = q.shape[1]
rep = torch.mean(v, dim=1, keepdim=True)
rep = rep.repeat(1, total_points, 1)
return rep
def _laplace_attention(self, k, v, q, scale=0.5):
k_ = k.unsqueeze(1)
v_ = v.unsqueeze(2)
unnorm_weights = torch.abs((k_ - v_) * scale)
unnorm_weights = unnorm_weights.sum(dim=-1)
weights = torch.softmax(unnorm_weights, dim=-1)
rep = torch.einsum("bik,bkj->bij", weights, v)
return rep
def _dot_attention(self, k, v, q):
scale = q.shape[-1] ** 0.5
unnorm_weights = torch.einsum("bjk,bik->bij", k, q) / scale
weights = torch.softmax(unnorm_weights, dim=-1)
rep = torch.einsum("bik,bkj->bij", weights, v)
return rep
def _pytorch_multihead_attention(self, k, v, q):
# Pytorch multiheaded attention takes inputs if diff order and permutation
o = self._mattn(q.permute(1, 0, 2), k.permute(1, 0, 2), v.permute(1, 0, 2))[0]
return o.permute(1, 0, 2)
class LatentEncoder(nn.Module):
def __init__(
self,
input_dim,
hidden_dim=32,
latent_dim=32,
n_heads=8,
self_attention_type="multihead",
n_encoder_layers=3,
min_std=0.1,
dropout=0,
):
super(LatentEncoder, self).__init__()
self._input_layer = NPBlockRelu2d(input_dim, hidden_dim, dropout)
self._encoder = nn.Sequential(
*[
NPBlockRelu2d(hidden_dim, hidden_dim, dropout)
for _ in range(n_encoder_layers)
]
)
self._self_attention = Attention(
hidden_dim, self_attention_type, n_heads=n_heads, dropout=dropout
)
self._penultimate_layer = block_relu(hidden_dim, hidden_dim, dropout)
self._mean = nn.Linear(hidden_dim, latent_dim)
self._log_var = nn.Linear(hidden_dim, latent_dim)
self.min_std = min_std
def forward(self, x, y):
encoder_input = torch.cat([x, y], dim=-1)
encoded = self._input_layer(encoder_input)
encoded = self._encoder(encoded)
attention_output = self._self_attention(encoded, encoded, encoded)
mean_repr = attention_output.mean(dim=1)
mean_repr = torch.relu(self._penultimate_layer(mean_repr))
mean = self._mean(mean_repr)
log_var = self._log_var(mean_repr)
# Clip it in the log domain, so it can only approach self.min_std, this helps aboid mode collaprse
log_var = F.softplus(log_var) + math.log(self.min_std)
sigma = torch.exp(0.5 * log_var)
dist = torch.distributions.Normal(mean, sigma)
return dist, log_var
class DeterministicEncoder(nn.Module):
def __init__(
self,
input_dim,
x_dim,
hidden_dim=32,
n_d_encoder_layers=3,
self_attention_type="multihead",
cross_attention_type="multihead",
dropout=0,
n_heads=8,
):
super(DeterministicEncoder, self).__init__()
self._input_layer = NPBlockRelu2d(input_dim, hidden_dim, dropout)
self._d_encoder = nn.Sequential(
*[
NPBlockRelu2d(hidden_dim, hidden_dim, dropout)
for _ in range(n_d_encoder_layers)
]
)
self._self_attention = Attention(
hidden_dim, self_attention_type, dropout=dropout, n_heads=n_heads
)
self._cross_attention = Attention(
hidden_dim, cross_attention_type, dropout=dropout, n_heads=n_heads
)
self._target_transform = nn.Linear(x_dim, hidden_dim)
self._context_transform = nn.Linear(x_dim, hidden_dim)
def forward(self, context_x, context_y, target_x):
d_encoder_input = torch.cat([context_x, context_y], dim=-1)
d_encoded = self._input_layer(d_encoder_input)
d_encoded = self._d_encoder(d_encoded)
attention_output = self._self_attention(d_encoded, d_encoded, d_encoded)
q = self._target_transform(target_x)
k = self._context_transform(context_x)
q = self._cross_attention(k, d_encoded, q)
return q
class Decoder(nn.Module):
def __init__(
self,
x_dim,
y_dim,
hidden_dim=32,
latent_dim=32,
n_decoder_layers=3,
min_std=0.1,
dropout=0,
):
super(Decoder, self).__init__()
self._target_transform = NPBlockRelu2d(x_dim, hidden_dim, dropout)
hidden_dim_2 = 2 * hidden_dim + latent_dim
self._decoder = nn.Sequential(
*[
NPBlockRelu2d(hidden_dim_2, hidden_dim_2, dropout)
for _ in range(n_decoder_layers)
]
)
self._mean = nn.Linear(2 * hidden_dim + latent_dim, y_dim)
self._std = nn.Linear(2 * hidden_dim + latent_dim, y_dim)
self.min_std = min_std
def forward(self, r, z, target_x):
x = self._target_transform(target_x)
representation = torch.cat([torch.cat([r, z], dim=-1), x], dim=-1)
representation = self._decoder(representation)
mean = self._mean(representation)
log_sigma = self._std(representation)
log_sigma = F.softplus(log_sigma) + math.log(self.min_std)
sigma = torch.exp(log_sigma)
dist = torch.distributions.Normal(mean, sigma)
return dist, log_sigma