<file_sep>from data_loader import *
from model import *
import pickle
import numpy as np

import torch
import torch.optim as optim
import matplotlib
import matplotlib.pyplot as plt
import time
import random

matplotlib.use('Agg')

def main(task, kind, n_epoch):
    loader = GPGenerator(batch_size=16, num_classes=1, data_source='gp', is_train=True)
    
    model = NPregression(kind=kind, dim_input=1, dim_output=1).to(device)
    
    optimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-4)
    
    KL_list, NLL_list = [], []
    train_loss_list, test_loss_list = [], []
    train_MSE_list, test_MSE_list = [], []
    for epoch in range(n_epoch):
        model.train()
        
        (Cx, Tx), (Cy, Ty), _, _ = loader.generate_mixture_batch(is_test=False)    
            
        Cx = torch.tensor(Cx, dtype=torch.float)
        Cy = torch.tensor(Cy, dtype=torch.float)
        Tx = torch.tensor(Tx, dtype=torch.float)
        Ty = torch.tensor(Ty, dtype=torch.float)
        
        _, train_KL, train_NLL, train_MSE = model(Cx.to(device), Cy.to(device), Tx.to(device), Ty.to(device))
        train_loss = train_KL + train_NLL
        
        optimizer.zero_grad()
        train_loss.backward()
        optimizer.step()
                    
        KL_list.append(train_KL.item())
        NLL_list.append(train_NLL.item())
        train_loss_list.append(train_loss.item())
        train_MSE_list.append(train_MSE.item())
        
        with torch.no_grad():
            model.eval()
        
            (Cx, Tx), (Cy, Ty), _, _ = loader.generate_mixture_batch(is_test=False)    
        
            Cx = torch.tensor(Cx, dtype=torch.float)
            Cy = torch.tensor(Cy, dtype=torch.float)
            Tx = torch.tensor(Tx, dtype=torch.float)
            Ty = torch.tensor(Ty, dtype=torch.float)
        
            _, _, test_NLL, test_MSE = model(Cx.to(device), Cy.to(device), Tx.to(device), Ty.to(device))
            test_loss = test_NLL
        
            test_loss_list.append(test_loss.item())
            test_MSE_list.append(test_MSE.item())
            
        if (epoch+1) % 10000 == 0:
            print('[Epoch %d] KL: %.3f, NLL: %.3f, train_loss: %.3f, test_loss: %.3f, train_MSE: %.3f, test_MSE: %.3f' 
                  % (epoch+1, train_KL, train_NLL, train_loss, test_loss, train_MSE, test_MSE))
            torch.save(model.state_dict(), 'models/'+task+kind+'/'+str(epoch+1)+'.pt')
            
            torch.save(KL_list, 'loss/'+task+kind+'/KL_list.pt')
            torch.save(NLL_list, 'loss/'+task+kind+'/NLL_list.pt')
            torch.save(train_loss_list, 'loss/'+task+kind+'/train_loss_list.pt')
            torch.save(test_loss_list, 'loss/'+task+kind+'/test_loss_list.pt')
            torch.save(train_MSE_list, 'loss/'+task+kind+'/train_MSE_list.pt')
            torch.save(test_MSE_list, 'loss/'+task+kind+'/test_MSE_list.pt')
        
            
def regression(task, kind):
    model = NPregression(kind=kind, dim_input=1, dim_output=1).to(device)
    model.load_state_dict(torch.load('models/'+task+kind+'/1000000.pt'))
    model.eval()
    result_path = 'plots/'+task+kind
    
    loader = GPGenerator(batch_size=2, num_classes=1, data_source='gp', is_train=False)
    
    (Cx, Tx), (Cy, Ty), _, _ = loader.generate_mixture_batch(is_test=True)    
    
    Cx = torch.tensor(Cx, dtype=torch.float)
    Cy = torch.tensor(Cy, dtype=torch.float)
    Tx = torch.tensor(Tx, dtype=torch.float)
    Ty = torch.tensor(Ty, dtype=torch.float)

    y_dist, _, _, MSE = model(Cx.to(device), Cy.to(device), Tx.to(device), Ty.to(device))
           
    Cx, Cy = torch.squeeze(Cx, -1).cpu(), torch.squeeze(Cy, -1).cpu()
    Tx, Ty = torch.squeeze(Tx, -1).cpu(), torch.squeeze(Ty, -1).cpu()
    
    mean, std = y_dist.mean.detach(), y_dist.stddev.detach()
    mean, std = torch.squeeze(mean, -1).cpu(), torch.squeeze(std, -1).cpu()
    
    plt.figure()
    plt.title('regression')
    plt.plot(Tx[0], Ty[0], 'k:', label='True')
    plt.plot(Cx[0], Cy[0], 'k^', markersize=10, label='Contexts')
    plt.plot(Tx[0], mean[0], 'b', label='Predictions')
    plt.fill(torch.cat((Tx[0], torch.flip(Tx[0], [0])),0),
             torch.cat((mean[0] - 1.96 * std[0], torch.flip(mean[0] + 1.96 * std[0], [0])),0),
             alpha=.5, fc='b', ec='None', label='95% confidence interval')
    plt.xlabel('$x$')
    plt.ylabel('$f(x)$')
    plt.ylim(min(Ty[0]) - 0.3 * (max(Ty[0]) - min(Ty[0])), max(Ty[0]) + 0.3 * (max(Ty[0]) - min(Ty[0])))
    plt.savefig(result_path+'.png')
    plt.close()

    
    
def evaluation(task, kind):
    model = NPregression(kind=kind, dim_input=1, dim_output=1).to(device)
    model.load_state_dict(torch.load('models/'+task+kind+'/1300000.pt'), strict=False)
    model.eval()
    
    print(1)
    loader = GPGenerator(batch_size=4000, num_classes=1, data_source='gp', is_train=False)
    (Cx, Tx), (Cy, Ty), _, _ = loader.generate_mixture_batch(is_test=True)    
        
    Cx = torch.tensor(Cx, dtype=torch.float)
    Cy = torch.tensor(Cy, dtype=torch.float)
    Tx = torch.tensor(Tx, dtype=torch.float)
    Ty = torch.tensor(Ty, dtype=torch.float)
        
    y_dist, _, _, MSE = model(Cx.to(device), Cy.to(device), Tx.to(device), Ty.to(device))
    
    NLL_tensor = -y_dist.log_prob(Ty.to(device)).sum(-1).mean(-1)
    
    mean = y_dist.mean.detach().cpu()        
    MSE_tensor = ((Ty-mean)**2).squeeze(-1).mean(-1)     
    print('NLL', NLL_tensor.mean(), 'MSE', MSE_tensor.mean())
    

def count_parameters():
    model = NPregression(kind=kind, dim_input=1, dim_output=1).to(device)
    return sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    
if __name__ == '__main__':  
    device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

    n_epoch = 1000000

    task = 'gp'
    kind = 'NP'
    
    tic = time.time()
    main(task, kind, n_epoch)
    toc = time.time()
    
    mon, sec = divmod(toc-tic, 60)
    hr, mon = divmod(mon, 60)
    print('total wall-clock time is ', int(hr),':',int(mon),':',int(sec))
<file_sep>import torch
import torch.nn as nn
import torch.nn.functional as F

import random
import math
import numpy as np
import matplotlib.pyplot as plt

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")     


class NPregression(nn.Module):
    def __init__(self, kind, dim_input, dim_output, dim_hidden=128, num_heads=8):
        super(NPRegression, self).__init__()
        self.kind = kind
        self.dim_input = dim_input
        self.dim_hidden = dim_hidden
        self.dim_output = dim_output
        
        if 'CNP' in self.kind:
            dec_dim = dim_input + dim_hidden
        else:
            dec_dim = dim_input + dim_hidden * 2
            
        self.x_MLP = nn.Sequential(nn.Linear(dim_input, dim_hidden),
                                   nn.ReLU(),
                                   nn.Linear(dim_hidden, dim_hidden))
        
        if 'LD' in self.kind:
            r_dim = dim_hidden * dim_output * 2
            z_dim = dim_hidden * dim_output * 2
            
            self.param_MLP = nn.Sequential(nn.Linear(dim_hidden, dim_hidden),
                                           nn.ReLU(),
                                           nn.Linear(dim_hidden, dim_hidden))
        
            self.layernorm = nn.LayerNorm(dim_hidden)
            
        else:
            r_dim = dim_hidden
            z_dim = dim_hidden
            self.decoding = nn.Sequential(nn.Linear(dec_dim, dim_hidden),
                                          nn.ReLU(),
                                          nn.Linear(dim_hidden, dim_hidden),
                                          nn.ReLU(),
                                          nn.Linear(dim_hidden, dim_hidden),
                                          nn.ReLU(),
                                          nn.Linear(dim_hidden, dim_hidden),
                                          nn.ReLU(),
                                          nn.Linear(dim_hidden, dim_output * 2))
    
        
        if 'ANP' in self.kind:
            self.e_MLP = nn.Sequential(nn.Linear(dim_input, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, r_dim))
        
            self.cross_Attention = nn.MultiheadAttention(r_dim, num_heads)
        
            self.r_Attention = nn.MultiheadAttention(r_dim, num_heads)
            
            self.sub_Layer = nn.Sequential(nn.Linear(r_dim, r_dim * 2),
                                           nn.ReLU(),
                                           nn.Linear(r_dim * 2, r_dim))
            
            self.layer_Norm = nn.LayerNorm(r_dim)
        
            self.s_Attention = nn.MultiheadAttention(dim_hidden, num_heads)
        
        self.r_MLP = nn.Sequential(nn.Linear(dim_hidden + dim_output, dim_hidden),
                                   nn.ReLU(),
                                   nn.Linear(dim_hidden, dim_hidden),
                                   nn.ReLU(),
                                   nn.Linear(dim_hidden, dim_hidden),
                                   nn.ReLU(),
                                   nn.Linear(dim_hidden, dim_hidden),
                                   nn.ReLU(),
                                   nn.Linear(dim_hidden, dim_hidden),
                                   nn.ReLU(),
                                   nn.Linear(dim_hidden, r_dim))
        
        self.s_MLP = nn.Sequential(nn.Linear(dim_hidden + dim_output, dim_hidden),
                                   nn.ReLU(),
                                   nn.Linear(dim_hidden, dim_hidden),
                                   nn.ReLU(),
                                   nn.Linear(dim_hidden, dim_hidden))
            
        self.z_MLP = nn.Sequential(nn.Linear(dim_hidden, dim_hidden),
                                   nn.ReLU(),
                                   nn.Linear(dim_hidden, z_dim * 2))
        
        
        self.sigmoid = nn.Sigmoid()
        self.softplus = nn.Softplus()
    
    def deterministic_encoder(self, det_enc_inp, Cx=None, Tx=None, attention=False, synthetic=True):
        if attention == False:
            ri = self.r_MLP(det_enc_inp)
            r = ri.mean(1, True)
        else:
            if synthetic:
                _, num_context, _ = det_enc_inp.size()
                query = self.e_MLP(Tx).permute(1,0,2)
                key = self.e_MLP(Cx).permute(1,0,2)
                value = self.r_MLP(det_enc_inp).permute(1,0,2)
                out, _ = self.cross_Attention(query, key, value)
                r = out.permute(1,0,2).unsqueeze(2)
            else:
                hidden1 = self.r_MLP(det_enc_inp).permute(1,0,2)
                hidden2, _ = self.r_Attention(hidden1, hidden1, hidden1)
                ri, _ = self.r_Attention(hidden2, hidden2, hidden2)
                
                query = self.e_MLP(Tx).permute(1,0,2)
                key = self.e_MLP(Cx).permute(1,0,2)
                out, _ = self.cross_Attention(query, key, ri)
                sub1 = out + query
                sub2 = self.layer_Norm(sub1.permute(1,0,2))
                r = self.layer_Norm(self.sub_Layer(sub2)+sub2).unsqueeze(2)
        return r        
        
    def reparameterize(self, dist):
        eps = torch.randn(size = dist.stddev.size()).to(device)
        sample = dist.mean + dist.stddev * eps
        return sample
    
    def stochastic_distribution(self, z_param):
        half = int(z_param.size(-1)/2)
        z_mu = z_param[:, :, :half]
        z_omega = z_param[:, :, half:]
        z_sigma = 0.1+0.9*self.sigmoid(z_omega)
        z_dist = torch.distributions.normal.Normal(z_mu, z_sigma)
        return z_dist
    
    def stochastic_encoder(self, lat_enc_inp, attention=False):
        if attention == False:
            si = self.s_MLP(lat_enc_inp)
            s = si.mean(1, True)
            z_param = self.z_MLP(s)
        else:
            hidden1 = self.s_MLP(lat_enc_inp).permute(1,0,2)
            hidden2, _ = self.s_Attention(hidden1, hidden1, hidden1)
            si, _ = self.s_Attention(hidden2, hidden2, hidden2)
            s = si.permute(1,0,2).mean(1, True)
            z_param = self.z_MLP(s)
        z_dist = self.stochastic_distribution(z_param)
        return z_dist
    
    def decoder(self, x, r, z, linear):
        if linear == False:
            r, z = r.squeeze(2), z.squeeze(2)
            dec_inp = torch.cat((x, r, z), 2)
            param = self.decoding(dec_inp)
            mu = param[:, :, :self.dim_output]
            omega = param[:, :, self.dim_output:]
            sigma = 0.1+0.9*self.softplus(omega)
            dist = torch.distributions.normal.Normal(mu, sigma)
        else:
            r = r.reshape(-1, r.size(1), self.dim_output * 2, self.dim_hidden)
            z = z.reshape(-1, z.size(1), self.dim_output * 2, self.dim_hidden)
            dec_param = r + self.param_MLP(z)
            dec_param = self.layernorm(dec_param).permute(0,1,3,2)
            
            param = torch.matmul(self.x_MLP(x).unsqueeze(2), dec_param).squeeze()
            mu = param[:,:,:self.dim_output]
            omega = param[:,:,self.dim_output:]
            sigma = 0.1+0.9*self.softplus(omega)
            dist = torch.distributions.normal.Normal(mu, sigma)
        return dist

    def forward(self, Cx, Cy, Tx, Ty, synthetic=True):
        _, num_context, _ = Cx.size()
        _, num_target, _ = Ty.size()
        
        context = torch.cat((self.x_MLP(Cx), Cy), 2)
        target = torch.cat((self.x_MLP(Tx), Ty), 2)
        
        NLL = torch.tensor([0.0]).to(device)
        
        if 'LD' in self.kind:
            linear = True
        else:
            linear = False
        
        if 'CNP' in self.kind:
            r_c = self.deterministic_encoder(context)
            r_c = r_c.unsqueeze(1).repeat(1, num_target, 1, 1)
            r_t = self.deterministic_encoder(target)
            r_t = r_t.unsqueeze(1).repeat(1, num_target, 1, 1)
            
            z = torch.empty(0).to(device)
            
            KL = torch.tensor([0.0]).to(device)
        
        elif 'ANP' in self.kind:
            r_c = self.deterministic_encoder(context, Cx, Tx, attention=True, synthetic=synthetic)
            r_t = self.deterministic_encoder(target, Tx, Tx, attention=True, synthetic=synthetic)

            z_prior = self.stochastic_encoder(context, attention=True)               
            z_posterior = self.stochastic_encoder(target, attention=True)      
            if self.training:
                z_sample = self.reparameterize(z_posterior)
            else:
                z_sample = z_prior.mean
            z = z_sample.unsqueeze(1).repeat(1, num_target, 1, 1)
            
            KL = torch.distributions.kl.kl_divergence(z_posterior, z_prior).sum(-1).mean()
            
        elif 'NP' in self.kind:
            r_c = self.deterministic_encoder(context)
            r_c = r_c.unsqueeze(1).repeat(1, num_target, 1, 1)
            r_t = self.deterministic_encoder(target)
            r_t = r_t.unsqueeze(1).repeat(1, num_target, 1, 1)
                
            z_prior = self.stochastic_encoder(context)
            z_posterior = self.stochastic_encoder(target)
            if self.training:
                z_sample = self.reparameterize(z_posterior)
            else:
                z_sample = z_prior.mean
            z = z_sample.unsqueeze(1).repeat(1, num_target, 1, 1)

            KL = torch.distributions.kl.kl_divergence(z_posterior, z_prior).sum(-1).mean()
                    
        y_dist = self.decoder(Tx, r_c, z, linear=linear)          
        NLL += -y_dist.log_prob(Ty).sum(-1).mean()   
        MSE = ((Ty-y_dist.mean)**2).sum(-1).mean()
        
        return y_dist, KL, NLL, MSE

    
class NPclassification(nn.Module):
    def __init__(self, kind, dim_input, dim_output, dim_hidden=128, num_heads=8, dp_rate=0.0):
        super(NPClassification, self).__init__()
        self.kind = kind
        self.dim_input = dim_input
        self.dim_hidden = dim_hidden
        self.dim_output = dim_output
        
        r_dim = dim_hidden
        z_dim = dim_hidden
        
        if 'MetaFun' in self.kind:
            dec_dim = dim_hidden
        elif 'CNP' in self.kind:
            dec_dim = dim_hidden * 6
        else:
            dec_dim = dim_hidden * 11
        
        if 'LD' in self.kind:
            self.param_MLP = nn.Sequential(nn.Linear(dim_hidden, dim_hidden),
                                           nn.ReLU(),
                                           nn.Linear(dim_hidden, dim_hidden))
        
            self.layernorm = nn.LayerNorm(dim_hidden)
            
        else:
            self.decoding = nn.Sequential(nn.Linear(dec_dim, dim_hidden),
                                          nn.ReLU(),
                                          nn.Linear(dim_hidden, dim_hidden),
                                          nn.ReLU(),
                                          nn.Linear(dim_hidden, dim_hidden),
                                          nn.ReLU(),
                                          nn.Linear(dim_hidden, dim_hidden),
                                          nn.ReLU(),
                                          nn.Linear(dim_hidden, dim_output))
    
            
        if dim_input == 84*84*3:
            self.x_emb = ImageEmbedding(dim_hidden, dp_rate)
            
        else:
            self.x_emb = nn.Sequential(nn.Dropout(p=dp_rate),
                                       nn.Linear(dim_input, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden))
        
        if 'ANP' in self.kind:
            self.cross_Attention = nn.MultiheadAttention(dim_hidden, num_heads)
        
            self.r_Attention = nn.MultiheadAttention(dim_hidden, num_heads)
            
            self.s_Attention = nn.MultiheadAttention(dim_hidden, num_heads)
        
        
        if 'MetaFun' in self.kind:
            self.l1 = nn.Parameter(torch.Tensor([1]*dim_hidden))
        
            self.m_MLP = nn.Sequential(nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden))
            
            self.u_t_MLP = nn.Sequential(nn.Linear(dim_hidden * 2, dim_hidden),
                                         nn.ReLU(),
                                         nn.Linear(dim_hidden, dim_hidden),
                                         nn.ReLU(),
                                         nn.Linear(dim_hidden, dim_hidden))
            
            self.u_f_MLP = nn.Sequential(nn.Linear(dim_hidden * 2, dim_hidden),
                                         nn.ReLU(),
                                         nn.Linear(dim_hidden, dim_hidden),
                                         nn.ReLU(),
                                         nn.Linear(dim_hidden, dim_hidden))
            
            self.w_MLP = nn.Sequential(nn.Linear(dec_dim, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden * 2))
        
        else:
            self.r_MLP = nn.Sequential(nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, r_dim))
        
        if 'ANP' in self.kind or 'MetaFun' in self.kind:
            self.e_MLP = nn.Sequential(nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden))
            
        if 'NP' in self.kind or 'ANP' in self.kind:
            self.s_MLP = nn.Sequential(nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, dim_hidden))
            
            self.z_MLP = nn.Sequential(nn.Linear(dim_hidden, dim_hidden),
                                       nn.ReLU(),
                                       nn.Linear(dim_hidden, z_dim * 2))
        
        self.sigmoid = nn.Sigmoid()
        self.softplus = nn.Softplus()
    
    def kernel(self, a, b):
        a = torch.matmul(torch.unsqueeze(a, 2), (1/self.l1).diag())
        b = torch.matmul(torch.unsqueeze(b, 1), (1/self.l1).diag())
        norm = torch.norm(a-b, dim=-1)
        norm_squared = torch.pow(norm, 2)
        kernel = torch.exp(-0.5 * norm_squared)
        return kernel
    
    def deterministic_encoder(self, x, y, t0=None, attention=False, iterative=False):
        if iterative:
            tmp = torch.zeros((x.size(0), x.size(1), self.dim_output, self.dim_hidden)).to(device)
            r = torch.zeros((x.size(0), t0.size(1), self.dim_output, self.dim_hidden)).to(device)
                            
            query = self.e_MLP(t0)
            key = self.e_MLP(x)
                
            for iter_idx in range(5): 
                m_k = self.m_MLP(tmp)
                m = m_k.mean(2, True).repeat(1, 1, self.dim_output, 1)
                u_t = self.u_t_MLP(torch.cat((m_k, m), 3))
                u_f = self.u_f_MLP(torch.cat((m_k, m), 3))
                    
                u = y.unsqueeze(3) * u_t + (1-y.unsqueeze(3)) * u_f
                weight_c = self.kernel(key, key)
                weight_t = self.kernel(query, key)
                for class_idx in range(self.dim_output):
                    delta_r_c = torch.matmul(weight_c, u[:, :, class_idx, :])
                    tmp[:, :, class_idx, :] -= 0.1 * delta_r_c                            
                    delta_r_t_ = torch.matmul(weight_t, u[:, :, class_idx, :])
                    r[:, :, class_idx, :] -= 0.1 * delta_r_t
        
        else:
            r = torch.tensor([]).to(device)
            for batch_idx in range(x.size(0)):
                class_param = torch.tensor([]).to(device)
                for class_idx in range(self.dim_output):
                    index = class_idx == y[batch_idx].max(-1).indices
                    inp = x[batch_idx][index].reshape(1, -1, self.dim_hidden)    
                    if attention == False:
                        ri = self.r_MLP(inp).mean(1, True)
                    else:
                        query = self.e_MLP(t0[batch_idx:batch_idx+1]).permute(1,0,2)
                        key = self.e_MLP(inp).permute(1,0,2)
                        value = self.r_MLP(inp).permute(1,0,2)
                        out, _ = self.cross_Attention(query, key, value)
                        ri = out.permute(1,0,2).unsqueeze(2)
                    class_param = torch.cat((class_param, ri), -2)
                r = torch.cat((r, class_param), 0)
        return r        
        
    def reparameterize(self, dist):
        eps = torch.randn(size = dist.stddev.size()).to(device)
        sample = dist.mean + dist.stddev * eps
        return sample
    
    def stochastic_distribution(self, z_param):
        z_mu = z_param[:, :, :self.dim_hidden]
        z_omega = z_param[:, :, self.dim_hidden:]
        z_sigma = 0.1+0.9*self.sigmoid(z_omega)
        z_dist = torch.distributions.normal.Normal(z_mu, z_sigma)
        return z_dist
    
    def stochastic_encoder(self, x, y, attention=False:
        z_param = torch.tensor([]).to(device)
        for batch_idx in range(x.size(0)):
            class_param = torch.tensor([]).to(device)
            for class_idx in range(self.dim_output):
                index = class_idx == y[batch_idx].max(-1).indices
                inp = x[batch_idx][index].reshape(1, -1, self.dim_hidden)
                if attention == False:
                    si = self.s_MLP(inp).mean(1, True)
                    zi_param = self.z_MLP(si)
                else:
                    hidden1 = self.s_MLP(inp).permute(1,0,2)
                    hidden2, _ = self.s_Attention(hidden1, hidden1, hidden1)
                    si, _ = self.s_Attention(hidden2, hidden2, hidden2)
                    si = si.permute(1,0,2).mean(1, True)
                    zi_param = self.z_MLP(si)
                class_param = torch.cat((class_param, zi_param), 1)
            z_param = torch.cat((z_param, class_param), 0)
        return z_param
            
    def decoder(self, x, r, z, iterative=False, linear=True):
        if iterative:
            w_param = self.w_MLP(r)
            w_mu = w_param[:, :, :, :self.dim_hidden]
            w_omega = w_param[:, :, :, self.dim_hidden:]
            w_sigma = 0.1+0.9*self.softplus(w_omega)
        
            w_dist = torch.distributions.normal.Normal(w_mu, w_sigma)
            w = self.reparameterize(w_dist).permute(0,1,3,2)
            logits = torch.matmul(x.unsqueeze(2), w).squeeze()
        
        elif linear:
            if z == None:
                w = r
            else:
                w = r + self.param_MLP(z)
            w = self.layernorm(w).permute(0,1,3,2)
            logits = torch.matmul(x.unsqueeze(2), w).squeeze()
        
        else:
            if z == None:
                dec_inp = torch.cat((x, r), 2)
            else:
                dec_inp = torch.cat((x, r, z), 2)
            logits = self.decoding(dec_inp)
            
        dist = torch.distributions.one_hot_categorical.OneHotCategorical(logits=logits)    
        return dist

    def forward(self, Cx, Cy, Tx, Ty):
        c0 = self.x_emb(Cx)
        t0 = self.x_emb(Tx)
        
        _, num_context, _ = c0.size()
        _, num_target, _ = t0.size()
        
        NLL = torch.tensor([0.0]).to(device)
        
        iterative=False
        linear = False
        
        if 'CNP' in self.kind:
            r_c = self.deterministic_encoder(c0, Cy)
            r_t = self.deterministic_encoder(t0, Ty)
            
            if 'LD' in self.kind:
                linear = True
                r_c = r_c.unsqueeze(1).repeat(1, num_target, 1, 1)
                r_t = r_t.unsqueeze(1).repeat(1, num_target, 1, 1)
                
            else:
                r_c = r_c.reshape(r_c.size(0), 1, -1)
                r_c = r_c.repeat(1, num_target, 1)
                r_t = r_t.reshape(r_t.size(0), 1, -1)
                r_t = r_t.repeat(1, num_target, 1)
            
            KL = torch.tensor([0.0]).to(device)
            
            z = None
            
        elif 'ANP' in self.kind:
            r_c = self.deterministic_encoder(c0, Cy, t0, attention=True)
            r_t = self.deterministic_encoder(t0, Ty, t0, attention=True)
            
            z_c_param = self.stochastic_encoder(c0, Cy, attention=True)             
            z_prior = self.stochastic_distribution(z_c_param)
            z_t_param = self.stochastic_encoder(t0, Ty, attention=True)      
            z_posterior = self.stochastic_distribution(z_t_param)
            
            if self.training:
                z_sample = self.reparameterize(z_posterior)
            else:
                z_sample = z_prior.mean
            
            if 'LD' in self.kind:
                linear = True
                z = z_sample.unsqueeze(1).repeat(1, num_target, 1, 1)
            else:
                r_c = r_c.reshape(r_c.size(0), num_target, -1)
                r_t = r_t.reshape(r_t.size(0), num_target, -1)
                
                z_sample = z_sample.reshape(z_sample.size(0), 1, -1)
                z = z_sample.repeat(1, num_target, 1)
            
            KL = torch.distributions.kl.kl_divergence(z_posterior, z_prior).sum(-1).mean()
                        
        elif 'NP' in self.kind:
            r_c = self.deterministic_encoder(c0, Cy)
            r_t = self.deterministic_encoder(t0, Ty)
            
            z_c_param = self.stochastic_encoder(c0, Cy)             
            z_prior = self.stochastic_distribution(z_c_param)
            z_t_param = self.stochastic_encoder(t0, Ty)      
            z_posterior = self.stochastic_distribution(z_t_param)
            
            if self.training:
                z_sample = self.reparameterize(z_posterior)
            else:
                z_sample = z_prior.mean
            
            if 'LD' in self.kind:
                linear = True
                r_c = r_c.unsqueeze(1).repeat(1, num_target, 1, 1)
                r_t = r_t.unsqueeze(1).repeat(1, num_target, 1, 1)
                z = z_sample.unsqueeze(1).repeat(1, num_target, 1, 1)
            else:
                r_c = r_c.reshape(r_c.size(0), 1, -1)
                r_c = r_c.repeat(1, num_target, 1)
                r_t = r_t.reshape(r_t.size(0), 1, -1)
                r_t = r_t.repeat(1, num_target, 1)
                
                z_sample = z_sample.reshape(z_sample.size(0), 1, -1)
                z = z_sample.repeat(1, num_target, 1)
                    
            KL = torch.distributions.kl.kl_divergence(z_posterior, z_prior).sum(-1).mean()
                                
        elif 'MetaFun' in self.kind:
            iterative=True
            r_c = self.deterministic_encoder(c0, Cy, t0, iterative=iterative)
            
            KL = torch.tensor([0.0]).to(device)
            
            z = None
            
        y_dist = self.decoder(t0, r_c, z, iterative, linear)          
        NLL += -y_dist.log_prob(Ty).sum(-1).mean()
        accur = ((Ty.max(-1).indices == y_dist.logits.max(-1).indices).sum(-1) * 100.0 / Ty.size(1)).mean()
        
        return y_dist, KL, NLL, accur
        

    
    
<file_sep># Following links are what we referred to for data generation
#https://github.com/deepmind/neural-processes

import numpy as np
import os
import random
from scipy import signal
from PIL import Image
import pdb
import itertools
import pickle
import torch 
import torchvision
import torch.nn.functional as F
import torchvision.transforms as Transforms


class GPGenerator(object):
    def __init__(self, batch_size, num_classes, data_source, is_train):
        super(GPGenerator, self).__init__()
        self.batch_size = batch_size
        self.num_classes = num_classes 
        
    def generate_mixture_batch(self, is_test):
        dim_input = 1
        dim_output = 1
        batch_size = self.batch_size
        update_batch_size = 8#np.random.randint(5, 11)
        update_batch_size_eval = 30 - update_batch_size
        num_samples_per_class = update_batch_size + update_batch_size_eval
        
        l = 0.5#np.random.uniform(0.1, 0.6, size=batch_size)
        sigma = 1.0#np.random.uniform(0.1, 1.0, size=batch_size)
        
        sel_set = np.zeros(batch_size)

        if is_test:
            init_inputs = np.zeros([batch_size, 1000, dim_input])
            outputs = np.zeros([batch_size, 1000, dim_output])
        else:
            init_inputs = np.zeros([batch_size, num_samples_per_class, dim_input])
            outputs = np.zeros([batch_size, num_samples_per_class, dim_output])

        for func in range(batch_size):
            if is_test:
                init_inputs[func] = np.expand_dims(np.linspace(-2.0, 2.0, num=1000), axis=1)
            else:
                init_inputs[func] = np.random.uniform(-2.0, 2.0,
                                                      size=(num_samples_per_class, dim_input))
                
            x1 = np.expand_dims(init_inputs[func], axis=0)
            x2 = np.expand_dims(init_inputs[func], axis=1)
            
            kernel = sigma**2 * np.exp(-0.5 * np.square(x1-x2) / l**2)
            kernel = np.sum(kernel, axis=-1)
            kernel += (0.02 ** 2) * np.identity(init_inputs[func].shape[0])
            cholesky = np.linalg.cholesky(kernel)
            outputs[func] = np.matmul(cholesky, np.random.normal(size=(init_inputs[func].shape[0], dim_output)))
            
        rng = np.random.default_rng()
        random_idx = rng.permutation(init_inputs.shape[1])
        
        Cx = init_inputs[:, random_idx[:update_batch_size], :]
        Cy = outputs[:, random_idx[:update_batch_size], :]
        Tx = init_inputs
        Ty = outputs
        
        funcs_params = {'l': l, 'sigma': sigma}
        return (Cx, Tx), (Cy, Ty), funcs_params, sel_set

    
class ImageCompletionData(object):
    def __init__(self, batch_size, data_source):
        super(ImageCompletionData, self).__init__()
        self.batch_size = batch_size
        self.data_source = data_source
        
    def loader_generator(self, is_train, is_test):  
        transform = Transforms.Compose([Transforms.Resize(size=(32, 32)), Transforms.ToTensor()])
        
        if self.data_source == 'MNIST':
            dataset = torchvision.datasets.MNIST(root='../data/MNIST', train=is_train,
                                                   download=True, transform=Transforms.ToTensor())
            loader = torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, shuffle=is_train,
                                                 num_workers=8, drop_last=True)

        elif self.data_source == 'CIFAR10':
            dataset = torchvision.datasets.CIFAR10(root='../data/CIFAR10', train=is_train,
                                                   download=True, transform=transform)
            loader = torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, shuffle=is_train,
                                                 num_workers=8, drop_last=True)

        elif self.data_source == 'CelebA':
            dataset = torchvision.datasets.ImageFolder(root='../data/CelebA', transform=transform)
            if is_train:
                data = torchvision.datasets.ImageFolder(root='../data/CelebA', transform=transform)
                data.samples = dataset.samples[:162770]
                data.targets = dataset.targets[:162770]
                loader = torch.utils.data.DataLoader(data, batch_size=self.batch_size, shuffle=True, num_workers=8, drop_last=True)
            else:
                data = torchvision.datasets.ImageFolder(root='../data/CelebA', transform=transform)
                if is_test:
                    data.samples = dataset.samples[182637:202599]
                    data.targets = dataset.targets[182637:202599]
                else:
                    data.samples = dataset.samples[162770:182637]
                    data.targets = dataset.targets[162770:182637]
                loader = torch.utils.data.DataLoader(data, batch_size=self.batch_size, shuffle=False, num_workers=8, drop_last=True)
        
        return loader
        
        
<file_sep># Neural Process Family

Pytorch implementation of neural process families ([CNP](https://arxiv.org/pdf/1807.01613.pdf), [NP](https://arxiv.org/pdf/1807.01622.pdf), [ANP](https://arxiv.org/pdf/1901.05761.pdf), [Meta-Fun](https://arxiv.org/pdf/1912.02738.pdf)) on both regression and classification.

## How the studies have evolved

Neural Process Family is devised to alternate the gaussian process with a scalable neural network.
Specifically, given a set of input-output pairs {Cx, Cy} and some input Tx, the model is expected to credibly estimate its corresponding output Ty.
Hence, it follows the problem setting of few-shot learning and exploits an encoder-decoder pipeline.
A permutation invariant set encoding r is first extracted from the {Cx, Cy}, and Tx is then feedforwarded to decoder along with r to estimate the parameters of the distribution of Ty.

Conditional Neural Process(CNP) was the first instantiation of the studies, which is trained to maximize the marginal likelihood.
Neural Process(NP) is a simple extension to CNP by incorporating the stochastic latent variable following Variational AutoEncoder(VAE).
Since the marginal likelihood is no more tractable, variational inference technique is applied and the model is trained by maximizing the Evidence Lower Bound(ELBO).
Based on the Kolmogorov Extension Theorem, NP is proven to be a stochastic process, however, many complex modules are required to avoid underfitting due to difficulty on approximate inference in nature of probabilistic models.

Attentive Neural Process(ANP) is one of the follow-up studies based on this idea such that multi-head attention and self-attention introduced in Transformer is used to consider the dependency between the set elements.
Meta-Fun further bridges to the functional gradient descent that implictly relaxes to the infinite dimensional representation space.

## Implementation

See `data_loader.py` for preparing the dataset for 1D gaussian process regression and 2D image completion.

See `model.py` for comparing how the models feedforward the set data. 
Although Neural Process Family is devised to mainly deal with the regression problems, we also expand to classification problems using the linear classifier proposed in several meta-learning algorithms such as [VERSA](https://arxiv.org/pdf/1805.09921.pdf), [LEO](https://arxiv.org/pdf/1807.05960.pdf), [CNAP](https://arxiv.org/pdf/1906.07697.pdf).
Specifically, a linear classifier for certain class label (e.g. 1) is constructed by processing a subset of the context input {Cx | Cy = 1} of the label.

<p align="center">
<img src='./plots/classification prediction.png' height="192">

See `main.py`, for the hyperparameter setting and the overall training process.

As a pilot study, the models are evaluated on Gaussian Process regression and CIFAR10 image completion. Followings are some qualitative analysis on performance.
(For a detailed empirical comparison between the neural process families, refer to [paper](http://bayesiandeeplearning.org/2018/papers/92.pdf).)

While ANP converges much faster than NP and better handles underfitting, its interpolation and extrapolation turns out to be wiggly.

<p align="center">
<img src='./plots/NP_1D_training.gif' height="256">, <img src='./plots/ANP_1D_training.gif' height="256">

The prediction gets accurate as the number of context dataset increases.
<img src='./plots/ANP_2D_training.gif'>


