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

improved code working

Browse Source
This commit is contained in:
wassname
2020-02-01 14:55:02 +08:00
parent 5dfcd6e54f
commit 375a8f79e4
8 changed files with 3651 additions and 384 deletions
Download Patch File Download Diff File Expand all files Collapse all files
anp_1d_regression-from_src.ipynb
+612
View File
File diff suppressed because one or more lines are too long
anp_1d_regression.ipynb
+1974 -20
View File
File diff suppressed because one or more lines are too long
smartmeters-anp-optuna.ipynb
+795 -67
View File
File diff suppressed because it is too large Load Diff
src/models/attention.py
-120
View File
@@ -1,120 +0,0 @@
import torch as t
import torch.nn as nn
import torch.nn.functional as F
import torch
import math
class Linear(nn.Module):
"""
Linear Module
"""
def __init__(self, in_dim, out_dim, bias=True, w_init='linear'):
"""
:param in_dim: dimension of input
:param out_dim: dimension of output
:param bias: boolean. if True, bias is included.
:param w_init: str. weight inits with xavier initialization.
"""
super(Linear, self).__init__()
self.linear_layer = nn.Linear(in_dim, out_dim, bias=bias)
nn.init.xavier_uniform_(
self.linear_layer.weight,
gain=nn.init.calculate_gain(w_init))
def forward(self, x):
return self.linear_layer(x)
class MultiheadAttention(nn.Module):
"""
Multihead attention mechanism (dot attention)
"""
def __init__(self, num_hidden_k):
"""
:param num_hidden_k: dimension of hidden
"""
super(MultiheadAttention, self).__init__()
self.num_hidden_k = num_hidden_k
self.attn_dropout = nn.Dropout(p=0.1)
def forward(self, key, value, query):
# Get attention score
attn = t.bmm(query, key.transpose(1, 2))
attn = attn / math.sqrt(self.num_hidden_k)
attn = t.softmax(attn, dim=-1)
# Dropout
attn = self.attn_dropout(attn)
# Get Context Vector
result = t.bmm(attn, value)
return result, attn
class Attention(nn.Module):
"""
Attention Network
"""
def __init__(self, num_hidden, h=4):
"""
:param num_hidden: dimension of hidden
:param h: num of heads
"""
super(Attention, self).__init__()
self.num_hidden = num_hidden
self.num_hidden_per_attn = num_hidden // h
self.h = h
self.key = Linear(num_hidden, num_hidden, bias=False)
self.value = Linear(num_hidden, num_hidden, bias=False)
self.query = Linear(num_hidden, num_hidden, bias=False)
self.multihead = MultiheadAttention(self.num_hidden_per_attn)
self.residual_dropout = nn.Dropout(p=0.1)
self.final_linear = Linear(num_hidden * 2, num_hidden)
self.layer_norm = nn.LayerNorm(num_hidden)
def forward(self, key, value, query):
batch_size = key.size(0)
seq_k = key.size(1)
seq_q = query.size(1)
residual = query
# Make multihead
key = self.key(key).view(batch_size, seq_k, self.h, self.num_hidden_per_attn)
value = self.value(value).view(batch_size, seq_k, self.h, self.num_hidden_per_attn)
query = self.query(query).view(batch_size, seq_q, self.h, self.num_hidden_per_attn)
key = key.permute(2, 0, 1, 3).contiguous().view(-1, seq_k, self.num_hidden_per_attn)
value = value.permute(2, 0, 1, 3).contiguous().view(-1, seq_k, self.num_hidden_per_attn)
query = query.permute(2, 0, 1, 3).contiguous().view(-1, seq_q, self.num_hidden_per_attn)
# Get context vector
result, attns = self.multihead(key, value, query)
# Concatenate all multihead context vector
result = result.view(self.h, batch_size, seq_q, self.num_hidden_per_attn)
result = result.permute(1, 2, 0, 3).contiguous().view(batch_size, seq_q, -1)
# Concatenate context vector with input (most important)
result = t.cat([residual, result], dim=-1)
# Final linear
result = self.final_linear(result)
# Residual dropout & connection
result = self.residual_dropout(result)
result = result + residual
# Layer normalization
result = self.layer_norm(result)
return result, attns
src/models/lightning_anp.py
+5 -2
View File
@@ -24,24 +24,27 @@ class LatentModelPL(pl.LightningModule):
def training_step(self, batch, batch_idx):
assert all(torch.isfinite(d).all() for d in batch)
context_x, context_y, target_x, target_y = batch
y_pred, kl, loss, y_std = self.forward(context_x, context_y, target_x, target_y)
y_pred, kl, loss, loss_mse, y_std = self.forward(context_x, context_y, target_x, target_y)
tensorboard_logs = {
"train/loss": loss,
"train/kl": kl.mean(),
"train/std": y_std.mean(),
"train/mse": loss_mse.mean(),
"train/mse": F.mse_loss(y_pred, target_y).mean(),
}
assert torch.isfinite(loss)
# print('device', next(self.model.parameters()).device)
return {"loss": loss, "log": tensorboard_logs}
def validation_step(self, batch, batch_idx):
assert all(torch.isfinite(d).all() for d in batch)
context_x, context_y, target_x, target_y = batch
y_pred, kl, loss, y_std = self.forward(context_x, context_y, target_x, target_y)
y_pred, kl, loss, loss_mse, y_std = self.forward(context_x, context_y, target_x, target_y)
tensorboard_logs = {
"val_loss": loss,
"val/kl": kl.mean(),
"val/mse": loss_mse.mean(),
"val/std": y_std.mean(),
"val/mse": F.mse_loss(y_pred, target_y).mean(),
}
src/models/model.py
+54 -58
View File
@@ -33,33 +33,30 @@ def kl_loss_var(prior_mu, log_var_prior, post_mu, log_var_post):
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,
attention_dropout=0,
min_std=0.1,
use_lvar=False,
use_deterministic_path=True,
attention_layers=2,
**kwargs
):
def __init__(self,
x_dim,
y_dim,
hidden_dim=32,
latent_dim=32,
latent_enc_self_attn_type="dot",
det_enc_self_attn_type="dot",
det_enc_cross_attn_type="dot",
n_latent_encoder_layers=3,
n_det_encoder_layers=3,
n_decoder_layers=3,
use_deterministic_path=True,
min_std=0.01,
dropout=0,
use_self_attn=False,
attention_dropout=0,
batchnorm=False,
use_lvar=False,
attention_layers=2,
**kwargs,
):
super().__init__()
self.use_lvar = use_lvar
self.use_deterministic_path = use_deterministic_path
super(LatentModel, self).__init__()
self._latent_encoder = LatentEncoder(
x_dim + y_dim,
@@ -67,25 +64,27 @@ class LatentModel(nn.Module):
latent_dim=latent_dim,
self_attention_type=latent_enc_self_attn_type,
n_encoder_layers=n_latent_encoder_layers,
attention_layers=attention_layers,
dropout=dropout,
use_self_attn=use_self_attn,
attention_dropout=attention_dropout,
n_heads=num_heads,
batchnorm=batchnorm,
min_std=min_std,
use_lvar=use_lvar,
attention_layers=attention_layers,
)
self._deterministic_encoder = DeterministicEncoder(
x_dim + y_dim,
x_dim,
input_dim=x_dim + y_dim,
x_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,
attention_dropout=attention_dropout,
n_heads=num_heads,
attention_layers=attention_layers,
use_self_attn=use_self_attn,
dropout=dropout,
batchnorm=batchnorm,
attention_dropout=attention_dropout,
)
self._decoder = Decoder(
@@ -93,59 +92,56 @@ class LatentModel(nn.Module):
y_dim,
hidden_dim=hidden_dim,
latent_dim=latent_dim,
n_decoder_layers=n_decoder_layers,
dropout=dropout,
batchnorm=batchnorm,
min_std=min_std,
use_lvar=use_lvar,
use_deterministic_path=use_deterministic_path
n_decoder_layers=n_decoder_layers,
use_deterministic_path=use_deterministic_path,
)
self._use_deterministic_path = use_deterministic_path
self._use_lvar = use_lvar
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)
if self.training:
z = dist_post.rsample()
else:
# instead of sampling, in test mode take the mean, this will make it more deterministic
z = dist_post.loc
if target_y is not None:
dist_post, log_var_post = self._latent_encoder(target_x,
target_y)
z = dist_post.loc
else:
z = dist_prior.loc
z = z.unsqueeze(1).repeat(1, num_targets, 1) # [B, T_target, H]
if self.use_deterministic_path:
r = self._deterministic_encoder(
context_x, context_y, target_x
) # [B, T_target, H]
if self._use_deterministic_path:
r = self._deterministic_encoder(context_x, context_y,
target_x) # [B, T_target, H]
else:
r = None
dist, log_sigma = self._decoder(r, z, target_x)
dist, log_sigma = self._decoder(r, z, target_x)
if target_y is not None:
if self.use_lvar:
# Log likelihood has shape (batch_size, num_target, y_dim).
# log_p = log_prob_sigma(target_y, dist.loc, log_sigma).mean(-1)
log_p = dist.log_prob(target_y).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)
if self._use_lvar:
log_p = log_prob_sigma(target_y, dist.loc, log_sigma).mean(-1) # [B, T_target, Y].mean(-1)
kl_loss = kl_loss_var(dist_prior.loc, log_var_prior,
dist_post.loc, log_var_post).mean(-1) # [B, R].mean(-1)
else:
log_p = dist.log_prob(target_y).mean(-1)
kl_loss = torch.distributions.kl_divergence(dist_post, dist_prior).mean(-1)
kl_loss = torch.distributions.kl_divergence(
dist_post, dist_prior).mean(-1)
kl_loss = kl_loss[:, None].expand(log_p.shape)
mse_loss = F.mse_loss(dist.loc, target_y)
loss = (kl_loss - log_p).mean()
else:
log_p = None
mse_loss = None
kl_loss = None
loss = None
mse_loss = None
y_pred = dist.rsample() if self.training else dist.loc
return y_pred, kl_loss, loss, dist.scale
return y_pred, kl_loss, loss, mse_loss, dist.scale
src/models/modules.py
+210 -116
View File
@@ -2,17 +2,21 @@ import torch
from torch import nn
import torch.nn.functional as F
import math
from .attention import Attention as PtAttention
import numpy as np
# from .attention import Attention as PtAttention
class NPBlockRelu2d(nn.Module):
"""Block for Neural Processes."""
def __init__(self, in_channels, out_channels, dropout=0, norm=True):
def __init__(
self, in_channels, out_channels, dropout=0, batchnorm=False, bias=False
):
super().__init__()
self.linear = nn.Linear(in_channels, out_channels)
self.linear = nn.Linear(in_channels, out_channels, bias=bias)
self.act = nn.ReLU()
self.dropout = nn.Dropout2d(dropout)
self.norm = nn.BatchNorm2d(out_channels) if norm else False
self.norm = nn.BatchNorm2d(out_channels) if batchnorm else False
def forward(self, x):
# x.shape is (Batch, Sequence, Channels)
@@ -30,18 +34,86 @@ class NPBlockRelu2d(nn.Module):
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 BatchMLP(nn.Module):
"""Apply MLP to the final axis of a 3D tensor (reusing already defined MLPs).
Args:
input: input tensor of shape [B,n,d_in].
output_sizes: An iterable containing the output sizes of the MLP as defined
in `basic.Linear`.
Returns:
tensor of shape [B,n,d_out] where d_out=output_size
"""
def __init__(
self, input_size, output_size, num_layers=2, dropout=0, batchnorm=False
):
super().__init__()
self.input_size = input_size
self.output_size = output_size
self.num_layers = num_layers
self.initial = NPBlockRelu2d(
input_size, output_size, dropout=dropout, batchnorm=batchnorm
)
self.encoder = nn.Sequential(
*[
NPBlockRelu2d(
output_size, output_size, dropout=dropout, batchnorm=batchnorm
)
for _ in range(num_layers - 2)
]
)
self.final = nn.Linear(output_size, output_size)
def forward(self, x):
x = self.initial(x)
x = self.encoder(x)
return self.final(x)
class AttnLinear(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.linear = nn.Linear(in_channels, out_channels, bias=False)
torch.nn.init.normal_(self.linear.weight, std=in_channels ** -0.5)
def forward(self, x):
x = self.linear(x)
return x
class Attention(nn.Module):
def __init__(self, hidden_dim, attention_type, attention_layers=1, n_heads=8, dropout=0):
def __init__(
self,
hidden_dim,
attention_type,
attention_layers=2,
n_heads=8,
x_dim=1,
rep="mlp",
dropout=0,
batchnorm=False,
):
super().__init__()
self._rep = rep
if self._rep == "mlp":
self.batch_mlp_k = BatchMLP(
x_dim,
hidden_dim,
attention_layers,
dropout=dropout,
batchnorm=batchnorm,
)
self.batch_mlp_q = BatchMLP(
x_dim,
hidden_dim,
attention_layers,
dropout=dropout,
batchnorm=batchnorm,
)
if attention_type == "uniform":
self._attention_func = self._uniform_attention
elif attention_type == "laplace":
@@ -49,21 +121,30 @@ class Attention(nn.Module):
elif attention_type == "dot":
self._attention_func = self._dot_attention
elif attention_type == "multihead":
self._mattn = nn.ModuleList([torch.nn.MultiheadAttention(
hidden_dim, n_heads, bias=False, dropout=dropout
) for _ in range(attention_layers)])
self._attention_func = self._pytorch_multihead_attention
self._W_k = nn.ModuleList(
[AttnLinear(hidden_dim, hidden_dim) for _ in range(n_heads)]
)
self._W_v = nn.ModuleList(
[AttnLinear(hidden_dim, hidden_dim) for _ in range(n_heads)]
)
self._W_q = nn.ModuleList(
[AttnLinear(hidden_dim, hidden_dim) for _ in range(n_heads)]
)
self._W = AttnLinear(n_heads * hidden_dim, hidden_dim)
self._attention_func = self._multihead_attention
self.n_heads = n_heads
elif attention_type == "ptmultihead":
self._mattn = nn.ModuleList([PtAttention(
hidden_dim, n_heads
) for _ in range(attention_layers)])
self._attention_func = self._ptmultihead_fn
self.n_heads = n_heads
self._W = torch.nn.MultiheadAttention(
hidden_dim, n_heads, bias=False, dropout=dropout
)
self._attention_func = self._pytorch_multihead_attention
else:
raise NotImplementedError
def forward(self, k, v, q):
if self._rep == "mlp":
k = self.batch_mlp_k(k)
q = self.batch_mlp_q(q)
rep = self._attention_func(k, v, q)
return rep
@@ -90,153 +171,162 @@ class Attention(nn.Module):
rep = torch.einsum("bik,bkj->bij", weights, v)
return rep
def _multihead_attention(self, k, v, q):
outs = []
for i in range(self.n_heads):
k_ = self._W_k[i](k)
v_ = self._W_v[i](v)
q_ = self._W_q[i](q)
out = self._dot_attention(k_, v_, q_)
outs.append(out)
outs = torch.stack(outs, dim=-1)
outs = outs.view(outs.shape[0], outs.shape[1], -1)
rep = self._W(outs)
return rep
def _pytorch_multihead_attention(self, k, v, q):
# Pytorch multiheaded attention takes inputs if diff order and permutation
q = q.permute(1, 0, 2)
k = k.permute(1, 0, 2)
v = v.permute(1, 0, 2)
for attention in self._mattn:
o = attention(q, k, v)[0]
q, k, v = o, o, o
o = self._W(q, k, v)[0]
return o.permute(1, 0, 2)
def _ptmultihead_fn(self, k, v, q):
for attention in self._mattn:
o = attention(k, v, q)[0]
# print(k.shape, v.shape, q.shape, o.shape)
q, k, v = o, o, o
return o
class LatentEncoder(nn.Module):
"""
Latent Encoder [For prior, posterior]
"""
def __init__(
self,
input_dim,
hidden_dim=32,
latent_dim=32,
n_heads=8,
self_attention_type="multihead",
self_attention_type="dot",
n_encoder_layers=3,
min_std=0.1,
min_std=0.01,
batchnorm=False,
dropout=0,
attention_dropout=0,
use_lvar=False,
use_self_attn=False,
attention_layers=2,
):
super().__init__()
self.use_lvar = use_lvar
self._input_layer = NPBlockRelu2d(input_dim, hidden_dim, dropout)
self._encoder = nn.Sequential(
*[
NPBlockRelu2d(hidden_dim, hidden_dim, dropout)
self._input_layer = nn.Linear(input_dim, hidden_dim)
self._encoder = nn.ModuleList(
[
NPBlockRelu2d(
hidden_dim, hidden_dim, batchnorm=batchnorm, dropout=dropout
)
for _ in range(n_encoder_layers)
]
)
self._self_attention = Attention(
hidden_dim, self_attention_type, n_heads=n_heads, dropout=attention_dropout, attention_layers=attention_layers
)
self._penultimate_layer = block_relu(hidden_dim, hidden_dim, dropout)
if use_self_attn:
self._self_attention = Attention(
hidden_dim,
self_attention_type,
attention_layers,
rep="identity",
dropout=attention_dropout,
)
self._penultimate_layer = nn.Linear(hidden_dim, hidden_dim)
self._mean = nn.Linear(hidden_dim, latent_dim)
self._log_var = nn.Linear(hidden_dim, latent_dim)
self.min_std = min_std
self._min_std = min_std
self._use_lvar = use_lvar
self._use_self_attn = use_self_attn
def forward(self, x, y):
"""Encodes the inputs into one representation.
Args:
x: Tensor of shape [B,observations,d_x]. For this 1D regression
task this corresponds to the x-values.
y: Tensor of shape [B,observations,d_y]. For this 1D regression
task this corresponds to the y-values.
Returns:
- A normal distribution over tensors of shape [B, num_latents]
- log_var
"""
# Concat location (x) and value (y) along the filter axes
encoder_input = torch.cat([x, y], dim=-1)
# Pass final axis through MLP
encoded = self._input_layer(encoder_input)
encoded = self._encoder(encoded)
for layer in self._encoder:
encoded = torch.relu(layer(encoded))
# Self-attention aggregator
attention_output = self._self_attention(encoded, encoded, encoded)
mean_repr = attention_output.mean(dim=1)
# Aggregator: take the mean over all points
if self._use_self_attn:
attention_output = self._self_attention(encoded, encoded, encoded)
mean_repr = attention_output.mean(dim=1)
else:
mean_repr = encoded.mean(dim=1)
# Have further MLP layers that map to the parameters of the Gaussian latent
mean_repr = self._penultimate_layer(mean_repr)
mean_repr = torch.relu(self._penultimate_layer(mean_repr))
# Then apply further linear layers to output latent mu and log sigma
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 collapase
if self.use_lvar:
log_var = log_var + math.log(self.min_std)
# Clip it in the log domain, so it can only approach self.min_std, this helps avoid mode collapase
# 2 ways, a better but untested way using the more stable log domain, and the way from the deepmind repo
if self._use_lvar:
log_var = torch.clamp(F.logsigmoid(log_var), np.log(self._min_std))
sigma = torch.exp(0.5 * log_var)
else:
sigma = self.min_std + (1 - self.min_std) * torch.sigmoid(log_var * 0.5)
sigma = self._min_std + (1 - self._min_std) * torch.sigmoid(log_var * 0.5)
dist = torch.distributions.Normal(mean, sigma)
return dist, log_var
class DeterministicEncoder(nn.Module):
"""
Deterministic Encoder [r]
"""
def __init__(
self,
input_dim,
x_dim,
hidden_dim=32,
n_d_encoder_layers=3,
self_attention_type="multihead",
cross_attention_type="multihead",
self_attention_type="dot",
cross_attention_type="dot",
use_self_attn=False,
attention_layers=2,
batchnorm=False,
dropout=0,
attention_dropout=0,
n_heads=8,
attention_layers=2,
):
super().__init__()
self._input_layer = NPBlockRelu2d(input_dim, hidden_dim, dropout)
self._d_encoder = nn.Sequential(
*[
NPBlockRelu2d(hidden_dim, hidden_dim, dropout)
self._use_self_attn = use_self_attn
self._input_layer = nn.Linear(input_dim, hidden_dim)
self._d_encoder = nn.ModuleList(
[
NPBlockRelu2d(
hidden_dim,
hidden_dim,
batchnorm=batchnorm,
dropout=attention_dropout,
)
for _ in range(n_d_encoder_layers)
]
)
self._self_attention = Attention(
hidden_dim, self_attention_type, dropout=attention_dropout, n_heads=n_heads, attention_layers=attention_layers
)
if use_self_attn:
self._self_attention = Attention(
hidden_dim,
self_attention_type,
attention_layers,
rep="identity",
dropout=attention_dropout,
)
self._cross_attention = Attention(
hidden_dim, cross_attention_type, dropout=attention_dropout, n_heads=n_heads, attention_layers=attention_layers
hidden_dim,
cross_attention_type,
x_dim=x_dim,
attention_layers=attention_layers,
)
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):
# concat context location (x), context value (y)
# Concatenate x and y along the filter axes
d_encoder_input = torch.cat([context_x, context_y], dim=-1)
# Pass final axis through MLP
d_encoded = self._input_layer(d_encoder_input)
d_encoded = self._d_encoder(d_encoded)
for layer in self._d_encoder:
d_encoded = torch.relu(layer(d_encoded))
# Apply self attention
d_encoded = self._self_attention(d_encoded, d_encoded, d_encoded)
if self._use_self_attn:
d_encoded = self._self_attention(d_encoded, d_encoded, d_encoded)
# query: target_x, key: context_x, value: d_encoded (representation of x)
k = self._context_transform(context_x)
q = self._target_transform(target_x)
# Apply attention
h = self._cross_attention(context_x, d_encoded, target_x)
# Cross Attention
r = self._cross_attention(k, d_encoded, q)
return r
return h
class Decoder(nn.Module):
@@ -247,51 +337,55 @@ class Decoder(nn.Module):
hidden_dim=32,
latent_dim=32,
n_decoder_layers=3,
min_std=0.1,
dropout=0,
use_deterministic_path=True,
min_std=0.01,
use_lvar=False,
use_deterministic_path=True
batchnorm=False,
dropout=0,
):
super().__init__()
self.use_lvar = use_lvar
self._target_transform = NPBlockRelu2d(x_dim, hidden_dim, dropout)
self.use_deterministic_path = use_deterministic_path
super(Decoder, self).__init__()
self._target_transform = nn.Linear(x_dim, hidden_dim)
if use_deterministic_path:
hidden_dim_2 = 2 * hidden_dim + latent_dim
else:
hidden_dim_2 = hidden_dim + latent_dim
self._decoder = nn.Sequential(
*[
NPBlockRelu2d(hidden_dim_2, hidden_dim_2, dropout)
self._decoder = nn.ModuleList(
[
NPBlockRelu2d(
hidden_dim_2, hidden_dim_2, batchnorm=batchnorm, dropout=dropout
)
for _ in range(n_decoder_layers)
]
)
self._mean = nn.Linear(hidden_dim_2, y_dim)
self._std = nn.Linear(hidden_dim_2, y_dim)
self.min_std = min_std
self._use_deterministic_path = use_deterministic_path
self._min_std = min_std
self._use_lvar = use_lvar
def forward(self, r, z, target_x):
# concatenate target_x and representation
x = self._target_transform(target_x)
# concatenate target_x and representation
if self.use_deterministic_path:
if self._use_deterministic_path:
z = torch.cat([r, z], dim=-1)
representation = torch.cat([z, x], dim=-1)
# Pass final axis through MLP
representation = self._decoder(representation)
for layer in self._decoder:
representation = torch.relu(layer(representation))
# Get the mean and the variance
mean = self._mean(representation)
log_sigma = self._std(representation)
# Bound the variance
if self.use_lvar:
log_sigma = log_sigma + math.log(self.min_std)
# Bound or clamp the variance
if self._use_lvar:
log_sigma = torch.clamp(log_sigma, math.log(self._min_std))
sigma = torch.exp(log_sigma)
else:
sigma = self.min_std + (1-self.min_std) * F.softplus(log_sigma)
sigma = self._min_std + (1 - self._min_std) * F.softplus(log_sigma)
# Dist
dist = torch.distributions.Normal(mean, sigma)
return dist, log_sigma
src/plot.py
+1 -1
View File
@@ -84,7 +84,7 @@ def plot_from_loader(
model.eval()
with torch.no_grad():
y_pred, kl, loss_test, y_std = model(context_x, context_y, target_x, target_y)
y_pred, kl, loss_test, loss_mse, y_std = model(context_x, context_y, target_x, target_y)
if plot:
plt.figure()