Refactoring

Author: makora9143

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+runs/
+.ipynb_checkpoints/

1d_demo.ipynb

@@ -0,0 +1 @@
+from .dataset import Synthetic1D

2d_demo.ipynb

@@ -2,6 +2,8 @@
 from torch.utils import data as tdata
 from gpytorch.utils.cholesky import psd_safe_cholesky
 
+from .kernels import eq_kernel, matern_kernel, periodic_kernel
+
 
 class Synthetic1D(tdata.Dataset):
  def __init__(self,
@@ -17,7 +19,15 @@ def __init__(self,
  self.x_dim = 1
  self.y_dim = 1
 
- self.kernel = kernel
+ if kernel == 'eq':
+ self.kernel = eq_kernel
+ elif kernel == 'matern':
+ self.kernel = matern_kernel
+ elif kernel == 'periodic':
+ self.kernel = periodic_kernel
+ else:
+ raise NotImplementedError('{} kernel is not implemented'.format(kernel))
+
  self.length_scale = length_scale
  self.output_scale = output_scale
 

convcnp.py

@@ -0,0 +1,2 @@
+from .convcnp1d import ConvCNP1d
+from .convcnp2d import ConvCNP2d

convcnp/dataset/__init__.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+from torch.distributions import MultivariateNormal
+
+from gpytorch.kernels import RBFKernel, ScaleKernel
+
+from ..modules import PowerFunction
+
+
+class ConvCNP1d(nn.Module):
+ def __init__(self, density=16):
+ super().__init__()
+
+ self.density = density
+
+ self.psi = ScaleKernel(RBFKernel())
+ self.phi = PowerFunction(K=1)
+
+ self.cnn = nn.Sequential(
+ nn.Conv1d(3, 16, 5, 1, 2),
+ nn.ReLU(),
+ nn.Conv1d(16, 32, 5, 1, 2),
+ nn.ReLU(),
+ nn.Conv1d(32, 16, 5, 1, 2),
+ nn.ReLU(),
+ nn.Conv1d(16, 2, 5, 1, 2)
+ )
+
+ def weights_init(m):
+ if isinstance(m, nn.Conv1d):
+ torch.nn.init.xavier_uniform_(m.weight)
+ torch.nn.init.zeros_(m.bias)
+ self.cnn.apply(weights_init)
+
+ self.pos = nn.Softplus()
+ self.psi_rho = ScaleKernel(RBFKernel())
+
+ def forward(self, xc: torch.Tensor, yc: torch.Tensor, xt: torch.Tensor):
+ tmp = torch.cat([xc, xt], 1)
+ lower, upper = tmp.min(), tmp.max()
+ num_t = int((self.density * (upper - lower)).item())
+ t = torch.linspace(start=lower, end=upper, steps=num_t).reshape(1, -1, 1).repeat(xc.size(0), 1, 1).to(xc.device)
+
+ h = self.psi(t, xc).matmul(self.phi(yc))
+ h0, h1 = h.split(1, -1)
+ h1 = h1.div(h0 + 1e-8)
+ h = torch.cat([h0, h1], -1)
+
+ rep = torch.cat([t, h], -1).transpose(-1, -2)
+ f = self.cnn(rep).transpose(-1, -2)
+ f_mu, f_sigma = f.split(1, -1)
+
+ mu = self.psi_rho(xt, t).matmul(f_mu)
+
+ sigma = self.psi_rho(xt, t).matmul(self.pos(f_sigma))
+ return MultivariateNormal(mu, scale_tril=sigma.diag_embed())