From e275979569ecfd0b3b004c599d088a4b949dc56f Mon Sep 17 00:00:00 2001 From: PaulJeha Date: Wed, 10 May 2023 19:25:59 +0200 Subject: [PATCH 1/9] fixing some small things for discrete data --- Model/Energy/EnergyForDistribution/poisson.py | 2 - .../abstract_trainer.py | 105 ++++++++++-------- .../trainer_self_normalized.py | 6 +- Model/Utils/plot_utils.py | 10 +- .../YAMLEBM_2D/self_normalized.yaml | 2 +- .../YAMLENERGY/energy_categorical.yaml | 2 +- 6 files changed, 75 insertions(+), 52 deletions(-) diff --git a/Model/Energy/EnergyForDistribution/poisson.py b/Model/Energy/EnergyForDistribution/poisson.py index 0d32404..233c680 100644 --- a/Model/Energy/EnergyForDistribution/poisson.py +++ b/Model/Energy/EnergyForDistribution/poisson.py @@ -55,6 +55,4 @@ def forward( Returns: Float[torch.Tensor, "batch_size 1"], E(x), the energy of the poisson distribution """ - # print(f"self.lambda_: {self.lambda_}") - # print(f"x: {x.shape}") return torch.lgamma(x + 1) - x * torch.log(self.lambda_) diff --git a/Model/Trainer/DistributionEstimation/abstract_trainer.py b/Model/Trainer/DistributionEstimation/abstract_trainer.py index f838203..1a41e3c 100644 --- a/Model/Trainer/DistributionEstimation/abstract_trainer.py +++ b/Model/Trainer/DistributionEstimation/abstract_trainer.py @@ -1,3 +1,4 @@ +import itertools import os import matplotlib.pyplot as plt @@ -5,18 +6,14 @@ import pytorch_lightning as pl import torch import yaml +from torch.distributions import categorical + +from Dataset.MissingDataDataset.DiscreteDataset import dic_discrete_dataset from ...Sampler import get_sampler from ...Utils.optimizer_getter import get_optimizer, get_scheduler -from ...Utils.plot_utils import plot_energy_2d, plot_images -from ...Utils.proposal_loss import log_prob_kl_loss, kl_loss, log_prob_loss -from ...Sampler import get_sampler -import numpy as np -import os -import matplotlib.pyplot as plt -import yaml -import itertools - +from ...Utils.plot_utils import plot_energy_2d, plot_images, print_discrete_params +from ...Utils.proposal_loss import kl_loss, log_prob_kl_loss, log_prob_loss class AbstractDistributionEstimation(pl.LightningModule): @@ -33,9 +30,11 @@ def __init__( self.args_dict = args_dict print("args_dict", args_dict) self.hparams.update(args_dict) - self.last_save = -float('inf') # To save the energy contour plot - self.last_save_sample = 0 # To save the samples - self.sampler = get_sampler(args_dict,) + self.last_save = -float("inf") # To save the energy contour plot + self.last_save_sample = 0 # To save the samples + self.sampler = get_sampler( + args_dict, + ) self.transform_back = complete_dataset.transform_back self.nb_sample_train_estimate = nb_sample_train_estimate @@ -237,17 +236,25 @@ def proposal_visualization(self): ) def configure_optimizers(self): - params_ebm = [child.parameters() for name,child in self.ebm.named_children() if name != 'proposal'] + params_ebm = [ + child.parameters() + for name, child in self.ebm.named_children() + if name != "proposal" + ] params_ebm.append(self.ebm.parameters()) - params_proposal = [self.ebm.proposal.parameters()] if self.ebm.proposal is not None else [] - ebm_opt = get_optimizer( args_dict = self.args_dict, list_params_gen = params_ebm) - proposal_opt = get_optimizer( args_dict = self.args_dict, list_params_gen = params_proposal) - - ebm_sch = get_scheduler(args_dict = self.args_dict, optim = ebm_opt) - proposal_sch = get_scheduler(args_dict = self.args_dict, optim = proposal_opt) - if ebm_sch is not None and proposal_sch is not None : - return [ebm_opt, proposal_opt], [ebm_sch, proposal_sch] - elif ebm_sch is not None : + params_proposal = ( + [self.ebm.proposal.parameters()] if self.ebm.proposal is not None else [] + ) + ebm_opt = get_optimizer(args_dict=self.args_dict, list_params_gen=params_ebm) + proposal_opt = get_optimizer( + args_dict=self.args_dict, list_params_gen=params_proposal + ) + + ebm_sch = get_scheduler(args_dict=self.args_dict, optim=ebm_opt) + proposal_sch = get_scheduler(args_dict=self.args_dict, optim=proposal_opt) + if ebm_sch is not None and proposal_sch is not None: + return [ebm_opt, proposal_opt], [ebm_sch, proposal_sch] + elif ebm_sch is not None: return [ebm_opt, proposal_opt], ebm_sch elif proposal_sch is not None: return [ebm_opt, proposal_opt], proposal_sch @@ -354,30 +361,40 @@ def plot_samples(self, num_samples=None): save_dir = os.path.join(save_dir, "samples_energy") if not os.path.exists(save_dir): os.makedirs(save_dir) - samples, init_samples = self.samples_mcmc(num_samples=num_samples) - if np.prod(self.args_dict["input_size"]) == 2: - samples = samples.flatten(1) - plot_energy_2d( - self, - save_dir=save_dir, - samples=[samples], - samples_title=["HMC samples"], - name="samples", - step=self.global_step, - ) - # elif len(self.args_dict["input_size"]) == 2 and self.args_dict["input_size"][0]==1: - # plot_energy_1d() - elif len(self.args_dict["input_size"]) == 3: - plot_images( - algo=self, + if self.args_dict["dataset_name"] in dic_discrete_dataset.keys(): + print_discrete_params(self) + else: + samples, init_samples = self.samples_mcmc(num_samples=num_samples) + self._sample_categorical( + num_samples=num_samples, save_dir=save_dir, - images=samples, name="samples", step=self.global_step, - init_samples=init_samples, - transform_back=self.transform_back, ) - else: - raise NotImplementedError - self.last_save_sample = self.global_step + + if np.prod(self.args_dict["input_size"]) == 2: + samples = samples.flatten(1) + plot_energy_2d( + self, + save_dir=save_dir, + samples=[samples], + samples_title=["HMC samples"], + name="samples", + step=self.global_step, + ) + # elif len(self.args_dict["input_size"]) == 2 and self.args_dict["input_size"][0]==1: + # plot_energy_1d() + elif len(self.args_dict["input_size"]) == 3: + plot_images( + algo=self, + save_dir=save_dir, + images=samples, + name="samples", + step=self.global_step, + init_samples=init_samples, + transform_back=self.transform_back, + ) + else: + raise NotImplementedError + self.last_save_sample = self.global_step diff --git a/Model/Trainer/DistributionEstimation/trainer_self_normalized.py b/Model/Trainer/DistributionEstimation/trainer_self_normalized.py index f53fa84..e5acdf0 100644 --- a/Model/Trainer/DistributionEstimation/trainer_self_normalized.py +++ b/Model/Trainer/DistributionEstimation/trainer_self_normalized.py @@ -11,6 +11,7 @@ from ...Utils.plot_utils import plot_energy_2d, plot_images from .abstract_trainer import AbstractDistributionEstimation + class SelfNormalizedTrainer(AbstractDistributionEstimation): """ Trainer for the an importance sampling estimator of the partition function, which can be either importance sampling (with log) or self.normalized (with exp). @@ -42,8 +43,8 @@ def training_step(self, batch, batch_idx): and self.global_step == self.args_dict["switch_mode"] ): self.ebm.switch_mode() - x = batch['data'] - if hasattr(self.ebm.proposal, 'set_x'): + x = batch["data"] + if hasattr(self.ebm.proposal, "set_x"): self.ebm.proposal.set_x(x) energy_samples, dic_output = self.ebm.calculate_energy(x) @@ -89,5 +90,4 @@ def training_step(self, batch, batch_idx): x, dic_output, ) - return loss_total diff --git a/Model/Utils/plot_utils.py b/Model/Utils/plot_utils.py index 96e2b94..6324584 100644 --- a/Model/Utils/plot_utils.py +++ b/Model/Utils/plot_utils.py @@ -3,9 +3,17 @@ import matplotlib.pyplot as plt import numpy as np import torch +import torch.nn.functional as F import torchvision -# def plot_energy_1d(algo, save_dir, samples = [], samples_title = [], step, energy_type=True): + +def print_discrete_params(algo): + if algo.args_dict["dataset_name"] == "poisson": + print(f"self.ebm.energy.theta {algo.ebm.energy.lambda_}") + elif algo.args_dict["dataset_name"] == "categorical": + print(f"self.ebm.energy.theta {F.softmax(algo.ebm.energy.theta)}") + else: + raise NotImplementedError def plot_energy_2d( diff --git a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml index bd472c8..ff1d3c9 100644 --- a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml +++ b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml @@ -1,5 +1,5 @@ ebm_name : 'self_normalized' -max_epoch : 75 +max_epoch : 130 max_steps : 2000 save_energy_every : 100 samples_every: 500 diff --git a/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_categorical.yaml b/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_categorical.yaml index 2d8f651..9eb1ae8 100644 --- a/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_categorical.yaml +++ b/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_categorical.yaml @@ -1,4 +1,4 @@ energy_name: 'categorical' -energy_param: +energy_params: theta: null learn_theta: true \ No newline at end of file From b4e34f62118d04e381d74649605aa531a4208cb8 Mon Sep 17 00:00:00 2001 From: PaulJeha Date: Thu, 11 May 2023 11:29:04 +0200 Subject: [PATCH 2/9] update ising to be Erdos Renyin like in oops --- Model/Energy/EnergyForDistribution/ising.py | 26 ++++++++++++++------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/Model/Energy/EnergyForDistribution/ising.py b/Model/Energy/EnergyForDistribution/ising.py index 4a114bb..608f5fe 100644 --- a/Model/Energy/EnergyForDistribution/ising.py +++ b/Model/Energy/EnergyForDistribution/ising.py @@ -7,7 +7,7 @@ from jaxtyping import Float -class EnergyIsing(nn.Module): +class ErdosRenyiEnergyIsing(nn.Module): """Implement the energy of an Ising model. C.f. Oops I took a gradient from Grathwohl et al. According to table 1 of the paper the energy is defined as: @@ -25,7 +25,8 @@ class EnergyIsing(nn.Module): Attributes: W: nn.Linear (input_size, hidden_dim), the parameters of the energy. - W is initialized with a Bernoulli distribution with p=0.5. + W is initialized as the adjacency matrix of an Erdos-Renyi graph with probability p=4/prod(input_size). + So each node has an average degree of 4. b: torch.Tensor of size (hidden_dim), the parameters of the energy. b is initialized as a tensor of ones. @@ -38,12 +39,21 @@ def __init__( learn_b: bool = True, ) -> None: super().__init__() - self.W = nn.parameter.Parameter( - torch.ones(prod(input_size), prod(input_size)) * 0.5, requires_grad=learn_W - ) - self.b = nn.parameter.Parameter( - torch.ones(prod(input_size)), requires_grad=learn_b - ) + N = prod(input_size) + assert N > 4, "input_size is too small" + p = 1 / (N - 1) + G = torch.rand(N, N) < p + G = torch.triu(G, diagonal=1) + G = (G + G.T) * 1.0 + + weights = torch.randn_like(G) * ( + (1.0 / (N * p)) ** 0.5 + ) # From Oops I took a gradient + weights = weights * (1 - torch.tril(torch.ones_like(weights))) + weights = weights + weights.t() + + self.W = nn.parameter.Parameter(G * weights, requires_grad=learn_W) + self.b = nn.parameter.Parameter(torch.ones(N), requires_grad=learn_b) def forward( self, x: Float[torch.Tensor, "batch_size *dim"] From 6e892a66b70efb783b5b19ec212e97bfc911fe5c Mon Sep 17 00:00:00 2001 From: PaulJeha Date: Thu, 11 May 2023 14:42:21 +0200 Subject: [PATCH 3/9] Change code to oops i took a gradient --- Model/Energy/EnergyForDistribution/ising.py | 49 ++++++++++++--------- 1 file changed, 29 insertions(+), 20 deletions(-) diff --git a/Model/Energy/EnergyForDistribution/ising.py b/Model/Energy/EnergyForDistribution/ising.py index 608f5fe..7533477 100644 --- a/Model/Energy/EnergyForDistribution/ising.py +++ b/Model/Energy/EnergyForDistribution/ising.py @@ -1,10 +1,13 @@ from math import prod from typing import Tuple +import igraph as ig +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from jaxtyping import Float +from torch.distributions import bernoulli class ErdosRenyiEnergyIsing(nn.Module): @@ -35,37 +38,43 @@ class ErdosRenyiEnergyIsing(nn.Module): def __init__( self, input_size: Tuple[int], - learn_W: bool = True, - learn_b: bool = True, + n_node: int, + average_degree: int = 4, + init_bias: float = 0.0, + learn_G: bool = False, + learn_bias: bool = False, ) -> None: - super().__init__() - N = prod(input_size) - assert N > 4, "input_size is too small" - p = 1 / (N - 1) - G = torch.rand(N, N) < p - G = torch.triu(G, diagonal=1) - G = (G + G.T) * 1.0 - - weights = torch.randn_like(G) * ( - (1.0 / (N * p)) ** 0.5 - ) # From Oops I took a gradient + super(ErdosRenyiEnergyIsing, self).__init__() + # Code from Oops I took a gradient + g = ig.Graph.Erdos_Renyi(n_node, float(average_degree) / float(n_node)) + A = np.asarray(g.get_adjacency().data) # g.get_sparse_adjacency() + A = torch.tensor(A).float() + weights = torch.randn_like(A) * ((1.0 / avg_degree) ** 0.5) weights = weights * (1 - torch.tril(torch.ones_like(weights))) weights = weights + weights.t() - self.W = nn.parameter.Parameter(G * weights, requires_grad=learn_W) - self.b = nn.parameter.Parameter(torch.ones(N), requires_grad=learn_b) + self.G = nn.Parameter(A * weights, requires_grad=learn_G) + self.bias = nn.Parameter( + torch.ones((n_node,)).float() * init_bias, requires_grad=learn_bias + ) + self.data_dim = n_node def forward( - self, x: Float[torch.Tensor, "batch_size *dim"] + self, x: Float[torch.Tensor, "batch_size nb_point_in_graph"] ) -> Float[torch.Tensor, "batch_size"]: """Compute the energy of the Ising model. Args: - x: Float[torch.Tensor, "batch_size *dim"], batch input of the energy. + x: Float[torch.Tensor, "batch_size nb_point_in_graph"], batch input of the energy. + x is a vector of zeros and ones. Returns: Float[torch.Tensor, "batch_size"], E(x), the energy of the Ising model. """ - x = x.flatten(1) - Wx = torch.matmul(x, self.W.T) - return -torch.sum(x * Wx, dim=1) - x @ self.b + # code from Oops I took a gradient + + x = 2 * x - 1 # convert 0/1 to -1/1 + xg = x @ self.G + xgx = (xg * x).sum(-1) + b = (self.bias[None, :] * x).sum(-1) + return -xgx - b From 21f36513ab1094fcee8a497d758ab3a89de9aeb2 Mon Sep 17 00:00:00 2001 From: PaulJeha Date: Fri, 12 May 2023 23:53:12 +0200 Subject: [PATCH 4/9] Coding Ising proposal --- Dataset/MissingDataDataset | 2 +- .../ising_proposal.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) create mode 100644 Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py diff --git a/Dataset/MissingDataDataset b/Dataset/MissingDataDataset index 9c3a5bd..2cfdf9f 160000 --- a/Dataset/MissingDataDataset +++ b/Dataset/MissingDataDataset @@ -1 +1 @@ -Subproject commit 9c3a5bd6ba0ac8dfd8cab9dc1d53ff3b6dc42e0f +Subproject commit 2cfdf9f7a9ed29ea6a2e160578070d684ecc87b3 diff --git a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py new file mode 100644 index 0000000..e0c7976 --- /dev/null +++ b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py @@ -0,0 +1,27 @@ +import numpy as np +import torch +import torch.nn as nn +from jaxtyping import Float + + +class IsingProposal(nn.Module): + def __init__(self, dataset): + super(IsingProposal, self).__init__() + self.dateset = dataset + + def sample( + self, nb_sample: int = 1 + ) -> Float[torch.Tensor, "nb_sample nb_point_in_graph"]: + if nb_sample < len(self.dateset): + index = np.random.choice(len(self.dateset), nb_sample) + + else: + index = np.random.choice(len(self.dateset), nb_sample, replace=True) + + center = torch.cat([self.dateset[i][0] for i in index]) + bernoulli_keep = torch.distributions.Bernoulli( + torch.full_like(center, 0.9) + ).sample() + samples = center * bernoulli_keep + (1 - center) * bernoulli_keep + + return samples.detach() From 33607d5ddd64752a8584aeaf961ed5b45b10449a Mon Sep 17 00:00:00 2001 From: PaulJeha Date: Sat, 13 May 2023 00:43:59 +0200 Subject: [PATCH 5/9] coding ising --- .../Energy/EnergyForDistribution/__init__.py | 2 +- Model/Energy/EnergyForDistribution/ising.py | 2 +- .../categorical.py | 2 +- .../ising_proposal.py | 40 +++++++-- Model/Proposals/proposal_getter.py | 81 ++++++++++++------- .../YAMLENERGY/energy_ising.yaml | 9 ++- .../YAMLDISTRIBUTION/YAMLPROPOSAL/ising.yaml | 3 + 7 files changed, 96 insertions(+), 43 deletions(-) create mode 100644 Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising.yaml diff --git a/Model/Energy/EnergyForDistribution/__init__.py b/Model/Energy/EnergyForDistribution/__init__.py index f223767..6c1ecb7 100644 --- a/Model/Energy/EnergyForDistribution/__init__.py +++ b/Model/Energy/EnergyForDistribution/__init__.py @@ -1,6 +1,6 @@ from .categorical import EnergyCategoricalDistrib from .conv import ConvEnergy -from .ising import EnergyIsing +from .ising import ErdosRenyiEnergyIsing from .linear import fc_energy from .poisson import EnergyPoissonDistribution from .rbm import EnergyRBM diff --git a/Model/Energy/EnergyForDistribution/ising.py b/Model/Energy/EnergyForDistribution/ising.py index 7533477..ef5b11b 100644 --- a/Model/Energy/EnergyForDistribution/ising.py +++ b/Model/Energy/EnergyForDistribution/ising.py @@ -49,7 +49,7 @@ def __init__( g = ig.Graph.Erdos_Renyi(n_node, float(average_degree) / float(n_node)) A = np.asarray(g.get_adjacency().data) # g.get_sparse_adjacency() A = torch.tensor(A).float() - weights = torch.randn_like(A) * ((1.0 / avg_degree) ** 0.5) + weights = torch.randn_like(A) * ((1.0 / average_degree) ** 0.5) weights = weights * (1 - torch.tril(torch.ones_like(weights))) weights = weights + weights.t() diff --git a/Model/Proposals/ProposalForDistributionEstimation/categorical.py b/Model/Proposals/ProposalForDistributionEstimation/categorical.py index 0c9785b..3f8de5b 100644 --- a/Model/Proposals/ProposalForDistributionEstimation/categorical.py +++ b/Model/Proposals/ProposalForDistributionEstimation/categorical.py @@ -28,7 +28,7 @@ def sample(self, nb_sample: int = 1): return samples_one_hot - def log_prob(self, x: Float[torch.Tensor, "batch_size"]): + def log_prob(self, x: Float[torch.Tensor, "batch_size 1"]): return ( categorical.Categorical(self.logit_parameters).log_prob(x).flatten(1).sum(1) ) diff --git a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py index e0c7976..90a8b92 100644 --- a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py +++ b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py @@ -5,23 +5,49 @@ class IsingProposal(nn.Module): - def __init__(self, dataset): + """Proposal for Ising model + + Attributes: + dataset: torch.utils.data.Dataset, dataset of the Ising model + p: float, probability of not flipping the spin + """ + + def __init__(self, dataset, p: float = 0.9): super(IsingProposal, self).__init__() self.dateset = dataset + self.p = p - def sample( - self, nb_sample: int = 1 - ) -> Float[torch.Tensor, "nb_sample nb_point_in_graph"]: + def sample(self, nb_sample: int = 1) -> Float[torch.Tensor, "nb_sample nb_nodes"]: if nb_sample < len(self.dateset): index = np.random.choice(len(self.dateset), nb_sample) - else: index = np.random.choice(len(self.dateset), nb_sample, replace=True) center = torch.cat([self.dateset[i][0] for i in index]) bernoulli_keep = torch.distributions.Bernoulli( - torch.full_like(center, 0.9) + torch.full_like(center, self.p) ).sample() - samples = center * bernoulli_keep + (1 - center) * bernoulli_keep + samples = center * bernoulli_keep + (1 - center) * (1 - bernoulli_keep) return samples.detach() + + def log_prob( + self, x: Float[torch.Tensor, "batch_size nb_nodes"] + ) -> Float[torch.Tensor, "batch_size"]: + data = torch.cat( + [self.dataset[i][0] for i in range(len(self.dataset))] + ).unsqueeze( + 1 + ) # len(dataset), 1, nb_nodes + x_expanded = x.unsqueeze(0) # 1, batch_size, nb_nodes + + dependency = (data - x_expanded).abs() # len(dataset), batch_size, nb_nodes + log_prob = ( + torch.distributions.Bernoulli(torch.full_like(dependency, self.p)) + .log_prob(dependency) + .sum(-1) + ) # len(dataset), batch_size + log_prob = log_prob.logsumexp(0) - torch.log( + torch.tensor(len(self.dateset)) + ) # batch_size + return log_prob diff --git a/Model/Proposals/proposal_getter.py b/Model/Proposals/proposal_getter.py index c954bf8..ee9d295 100644 --- a/Model/Proposals/proposal_getter.py +++ b/Model/Proposals/proposal_getter.py @@ -1,16 +1,17 @@ -from .ProposalForDistributionEstimation.standard_gaussian import StandardGaussian -from .ProposalForDistributionEstimation.kde import KernelDensity -from .ProposalForDistributionEstimation.kde_adaptive import KernelDensityAdaptive +import copy + +from .ProposalForDistributionEstimation.categorical import Categorical from .ProposalForDistributionEstimation.gaussian_mixture import GaussianMixtureProposal +from .ProposalForDistributionEstimation.gaussian_mixture_adaptive import ( + GaussianMixtureAdaptiveProposal, +) +from .ProposalForDistributionEstimation.ising_proposal import IsingProposal from .ProposalForDistributionEstimation.kde import KernelDensity +from .ProposalForDistributionEstimation.kde_adaptive import KernelDensityAdaptive from .ProposalForDistributionEstimation.poisson import Poisson from .ProposalForDistributionEstimation.standard_gaussian import StandardGaussian from .ProposalForRegression.MDNProposal import MDNProposalRegression from .ProposalForRegression.standard_gaussian import StandardGaussianRegression -from .ProposalForDistributionEstimation.categorical import Categorical -from .ProposalForDistributionEstimation.gaussian_mixture_adaptive import GaussianMixtureAdaptiveProposal - -import copy dic_proposals = { "standard_gaussian": StandardGaussian, @@ -18,49 +19,69 @@ "gaussian_mixture": GaussianMixtureProposal, "poisson": Poisson, "uniform_categorical": Categorical, - 'kernel_density_adaptive': KernelDensityAdaptive, - 'gaussian_mixture_adaptive': GaussianMixtureAdaptiveProposal, + "kernel_density_adaptive": KernelDensityAdaptive, + "gaussian_mixture_adaptive": GaussianMixtureAdaptiveProposal, + "ising": IsingProposal, } - def get_proposal(args_dict, input_size, dataset): - proposal = dic_proposals[args_dict["proposal_name"]] - if 'adaptive' in args_dict['proposal_name'] : - assert 'default_proposal_name' in args_dict.keys(), 'You need to specify a default proposal for the adaptive proposal' - assert not args_dict['train_proposal'], 'You cannot train the proposal if it is adaptive' - - if 'proposal_params' not in args_dict.keys(): - args_dict['proposal_params'] = {} + if "adaptive" in args_dict["proposal_name"]: + assert ( + "default_proposal_name" in args_dict.keys() + ), "You need to specify a default proposal for the adaptive proposal" + assert not args_dict[ + "train_proposal" + ], "You cannot train the proposal if it is adaptive" + + if "proposal_params" not in args_dict.keys(): + args_dict["proposal_params"] = {} aux_args_dict = copy.deepcopy(args_dict) - if 'default_proposal_params' not in args_dict.keys(): - aux_args_dict['proposal_params'] = {} - else : - aux_args_dict['proposal_params'] = args_dict['default_proposal_params'] - - aux_args_dict['proposal_name'] = args_dict['default_proposal_name'] - default_proposal = get_proposal(aux_args_dict, input_size, dataset,) - return proposal(default_proposal = default_proposal, input_size = input_size, dataset = dataset, **args_dict["proposal_params"]) + if "default_proposal_params" not in args_dict.keys(): + aux_args_dict["proposal_params"] = {} + else: + aux_args_dict["proposal_params"] = args_dict["default_proposal_params"] + + aux_args_dict["proposal_name"] = args_dict["default_proposal_name"] + default_proposal = get_proposal( + aux_args_dict, + input_size, + dataset, + ) + return proposal( + default_proposal=default_proposal, + input_size=input_size, + dataset=dataset, + **args_dict["proposal_params"] + ) if "proposal_params" in args_dict: return proposal(input_size, dataset, **args_dict["proposal_params"]) else: return proposal(input_size, dataset) + from .ProposalForRegression import UniformRegression dic_proposals_regression = { - 'standard_gaussian': StandardGaussianRegression, - 'mdn': MDNProposalRegression, - 'uniform': UniformRegression, + "standard_gaussian": StandardGaussianRegression, + "mdn": MDNProposalRegression, + "uniform": UniformRegression, } -def get_proposal_regression(args_dict, input_size_x, input_size_y, dataset,): +def get_proposal_regression( + args_dict, + input_size_x, + input_size_y, + dataset, +): proposal = dic_proposals_regression[args_dict["proposal_name"]] if "proposal_params" in args_dict: - proposal = proposal(input_size_x, input_size_y, dataset, **args_dict["proposal_params"]) + proposal = proposal( + input_size_x, input_size_y, dataset, **args_dict["proposal_params"] + ) return proposal else: return proposal(input_size_x, input_size_y, dataset) diff --git a/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml b/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml index e32b55f..60732e8 100644 --- a/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml +++ b/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml @@ -1,4 +1,7 @@ energy_name: 'ising' -energy_param: - learn_W: true - learn_b: true \ No newline at end of file +energy_params: + n_node: 100 + average_degree: 4 + init_bias: 0. + learn_G: true + learn_bias: true \ No newline at end of file diff --git a/Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising.yaml b/Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising.yaml new file mode 100644 index 0000000..e74783a --- /dev/null +++ b/Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising.yaml @@ -0,0 +1,3 @@ +proposal_name: "ising" +proposal_params: + p: 0.9 \ No newline at end of file From 62136f75774f0f3655d1e478cc2efa39c6457642 Mon Sep 17 00:00:00 2001 From: PaulJeha Date: Sat, 13 May 2023 01:09:32 +0200 Subject: [PATCH 6/9] DS STORE in gitignore --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index c5f5d86..4182f1c 100644 --- a/.gitignore +++ b/.gitignore @@ -129,4 +129,5 @@ dmypy.json .pyre/ /Results* -/Dataset/Downloaded/* \ No newline at end of file +/Dataset/Downloaded/* +.DS_STORE \ No newline at end of file From c25f186598984dd134e9ed042c1563ef46076214 Mon Sep 17 00:00:00 2001 From: PaulJeha Date: Mon, 15 May 2023 19:11:47 +0200 Subject: [PATCH 7/9] add ising --- Dataset/MissingDataDataset | 2 +- Model/Energy/EnergyForDistribution/ising.py | 13 ++- Model/Energy/energy_getter.py | 81 +++++++++++-------- .../ising_proposal.py | 18 +++-- .../abstract_trainer.py | 7 +- .../trainer_self_normalized.py | 11 +++ Model/Utils/plot_utils.py | 28 ++++++- .../YAMLEBM_2D/self_normalized.yaml | 8 +- 8 files changed, 117 insertions(+), 51 deletions(-) diff --git a/Dataset/MissingDataDataset b/Dataset/MissingDataDataset index 2cfdf9f..4791d39 160000 --- a/Dataset/MissingDataDataset +++ b/Dataset/MissingDataDataset @@ -1 +1 @@ -Subproject commit 2cfdf9f7a9ed29ea6a2e160578070d684ecc87b3 +Subproject commit 4791d397ce12ef0d90cf07493b1d436da2b12607 diff --git a/Model/Energy/EnergyForDistribution/ising.py b/Model/Energy/EnergyForDistribution/ising.py index ef5b11b..e2a3631 100644 --- a/Model/Energy/EnergyForDistribution/ising.py +++ b/Model/Energy/EnergyForDistribution/ising.py @@ -46,9 +46,10 @@ def __init__( ) -> None: super(ErdosRenyiEnergyIsing, self).__init__() # Code from Oops I took a gradient - g = ig.Graph.Erdos_Renyi(n_node, float(average_degree) / float(n_node)) - A = np.asarray(g.get_adjacency().data) # g.get_sparse_adjacency() - A = torch.tensor(A).float() + # g = ig.Graph.Erdos_Renyi(n_node, float(average_degree) / float(n_node)) + # A = np.asarray(g.get_adjacency().data) # g.get_sparse_adjacency() + # A = torch.tensor(A).float() + A = torch.randn((n_node, n_node)) * 0.01 weights = torch.randn_like(A) * ((1.0 / average_degree) ** 0.5) weights = weights * (1 - torch.tril(torch.ones_like(weights))) weights = weights + weights.t() @@ -59,6 +60,10 @@ def __init__( ) self.data_dim = n_node + @property + def J(self): + return self.G + def forward( self, x: Float[torch.Tensor, "batch_size nb_point_in_graph"] ) -> Float[torch.Tensor, "batch_size"]: @@ -74,7 +79,7 @@ def forward( # code from Oops I took a gradient x = 2 * x - 1 # convert 0/1 to -1/1 - xg = x @ self.G + xg = x @ self.J xgx = (xg * x).sum(-1) b = (self.bias[None, :] * x).sum(-1) return -xgx - b diff --git a/Model/Energy/energy_getter.py b/Model/Energy/energy_getter.py index ddb4def..18b3b9e 100644 --- a/Model/Energy/energy_getter.py +++ b/Model/Energy/energy_getter.py @@ -1,27 +1,31 @@ +import numpy as np + from .EnergyForDistribution import ( ConvEnergy, EnergyCategoricalDistrib, - EnergyIsing, EnergyPoissonDistribution, EnergyRBM, + ErdosRenyiEnergyIsing, fc_energy, ) -from .EnergyForRegression import EnergyNetworkRegression_Large, EnergyNetworkRegression_Toy +from .EnergyForRegression import ( + EnergyNetworkRegression_Large, + EnergyNetworkRegression_Toy, +) from .FeatureExtractor import Resnet18_FeatureExtractor, ToyFeatureNet -import numpy as np dic_energy = { "fc": fc_energy, "conv": ConvEnergy, "rbm": EnergyRBM, "categorical": EnergyCategoricalDistrib, "poisson": EnergyPoissonDistribution, - "ising": EnergyIsing, + "ising": ErdosRenyiEnergyIsing, } dic_energy_regression = { - 'fc': EnergyNetworkRegression_Large, - 'toy': EnergyNetworkRegression_Toy, + "fc": EnergyNetworkRegression_Large, + "toy": EnergyNetworkRegression_Toy, } @@ -56,8 +60,8 @@ def get_energy_regression(input_size_x, input_size_y, args_dict): dic_feature_extractor = { - 'resnet' : Resnet18_FeatureExtractor, - 'toy' : ToyFeatureNet, + "resnet": Resnet18_FeatureExtractor, + "toy": ToyFeatureNet, } @@ -69,37 +73,48 @@ def get_feature_extractor( return None if args_dict["feature_extractor_name"] not in dic_feature_extractor: raise ValueError("Feature extractor name not valid") - - feature_extractor = dic_feature_extractor[args_dict['feature_extractor_name']] - if 'feature_extractor_params' not in args_dict.keys(): - args_dict['feature_extractor_params'] = {} - feature_extractor = feature_extractor(input_dim=input_size_x, **args_dict['feature_extractor_params']) - print(args_dict['train_feature_extractor']) - if args_dict['train_feature_extractor'] == False : + + feature_extractor = dic_feature_extractor[args_dict["feature_extractor_name"]] + if "feature_extractor_params" not in args_dict.keys(): + args_dict["feature_extractor_params"] = {} + feature_extractor = feature_extractor( + input_dim=input_size_x, **args_dict["feature_extractor_params"] + ) + print(args_dict["train_feature_extractor"]) + if args_dict["train_feature_extractor"] == False: for param in feature_extractor.parameters(): param.requires_grad = False return feature_extractor + from .ExplicitBiasForRegression import Layer1FC, Layer2FC, Layer3FC + dic_explicit_bias_regression = { - '1_layer_fc' : Layer1FC, - '2_layer_fc' : Layer2FC, - '3_layer_fc' : Layer3FC, - 'none' : None, + "1_layer_fc": Layer1FC, + "2_layer_fc": Layer2FC, + "3_layer_fc": Layer3FC, + "none": None, } -def get_explicit_bias_regression(args_dict, - input_size_x, - ): - if 'explicit_bias_name' not in args_dict: + +def get_explicit_bias_regression( + args_dict, + input_size_x, +): + if "explicit_bias_name" not in args_dict: return None - if args_dict['explicit_bias_name'] not in dic_explicit_bias_regression: - raise ValueError('Explicit bias name not valid') - if args_dict['explicit_bias_name'] is None or args_dict['explicit_bias_name'] == 'none' : - return None - explicit_bias = dic_explicit_bias_regression[args_dict['explicit_bias_name']] - if 'explicit_bias_params' not in args_dict.keys(): - args_dict['explicit_bias_params'] = {} - - explicit_bias = explicit_bias(input_size_x=input_size_x, **args_dict['explicit_bias_params']) - return explicit_bias \ No newline at end of file + if args_dict["explicit_bias_name"] not in dic_explicit_bias_regression: + raise ValueError("Explicit bias name not valid") + if ( + args_dict["explicit_bias_name"] is None + or args_dict["explicit_bias_name"] == "none" + ): + return None + explicit_bias = dic_explicit_bias_regression[args_dict["explicit_bias_name"]] + if "explicit_bias_params" not in args_dict.keys(): + args_dict["explicit_bias_params"] = {} + + explicit_bias = explicit_bias( + input_size_x=input_size_x, **args_dict["explicit_bias_params"] + ) + return explicit_bias diff --git a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py index 90a8b92..f444a39 100644 --- a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py +++ b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn from jaxtyping import Float +from torch.nn.parameter import Parameter class IsingProposal(nn.Module): @@ -12,18 +13,19 @@ class IsingProposal(nn.Module): p: float, probability of not flipping the spin """ - def __init__(self, dataset, p: float = 0.9): + def __init__(self, input_size, dataset, p: float = 0.9): super(IsingProposal, self).__init__() - self.dateset = dataset + self.dataset = dataset self.p = p + self.dummy_param = Parameter(torch.Tensor([0.0]), requires_grad=False) def sample(self, nb_sample: int = 1) -> Float[torch.Tensor, "nb_sample nb_nodes"]: - if nb_sample < len(self.dateset): - index = np.random.choice(len(self.dateset), nb_sample) + if nb_sample < len(self.dataset): + index = np.random.choice(len(self.dataset), nb_sample) else: - index = np.random.choice(len(self.dateset), nb_sample, replace=True) + index = np.random.choice(len(self.dataset), nb_sample, replace=True) - center = torch.cat([self.dateset[i][0] for i in index]) + center = torch.stack([self.dataset[i][0] for i in index]) bernoulli_keep = torch.distributions.Bernoulli( torch.full_like(center, self.p) ).sample() @@ -34,7 +36,7 @@ def sample(self, nb_sample: int = 1) -> Float[torch.Tensor, "nb_sample nb_nodes" def log_prob( self, x: Float[torch.Tensor, "batch_size nb_nodes"] ) -> Float[torch.Tensor, "batch_size"]: - data = torch.cat( + data = torch.stack( [self.dataset[i][0] for i in range(len(self.dataset))] ).unsqueeze( 1 @@ -48,6 +50,6 @@ def log_prob( .sum(-1) ) # len(dataset), batch_size log_prob = log_prob.logsumexp(0) - torch.log( - torch.tensor(len(self.dateset)) + torch.tensor(len(self.dataset)) ) # batch_size return log_prob diff --git a/Model/Trainer/DistributionEstimation/abstract_trainer.py b/Model/Trainer/DistributionEstimation/abstract_trainer.py index 1a41e3c..bc0d484 100644 --- a/Model/Trainer/DistributionEstimation/abstract_trainer.py +++ b/Model/Trainer/DistributionEstimation/abstract_trainer.py @@ -96,9 +96,14 @@ def post_train_step_handler(self, x, dic_output): def validation_step(self, batch, batch_idx): x = batch["data"] energy_samples, dic_output = self.ebm.calculate_energy(x) + return dic_output def validation_epoch_end(self, outputs): + if self.args_dict["dataset_name"] == "ising": + with torch.no_grad(): + log_rmse_val = (self.J - self.ebm.energy.J).pow(2).mean().sqrt().log() + self.log_rmse_val.append((log_rmse_val.item(), self.global_step)) self.update_dic_logger(outputs, name="val_") self.proposal_visualization() self.plot_energy() @@ -363,7 +368,7 @@ def plot_samples(self, num_samples=None): os.makedirs(save_dir) if self.args_dict["dataset_name"] in dic_discrete_dataset.keys(): - print_discrete_params(self) + print_discrete_params(self, save_dir=save_dir, step=self.global_step) else: samples, init_samples = self.samples_mcmc(num_samples=num_samples) self._sample_categorical( diff --git a/Model/Trainer/DistributionEstimation/trainer_self_normalized.py b/Model/Trainer/DistributionEstimation/trainer_self_normalized.py index e5acdf0..797485a 100644 --- a/Model/Trainer/DistributionEstimation/trainer_self_normalized.py +++ b/Model/Trainer/DistributionEstimation/trainer_self_normalized.py @@ -34,6 +34,12 @@ def __init__( **kwargs, ) + if args_dict["dataset_name"] == "ising": + self.J = complete_dataset.J + self.log_rmse = [] + if args_dict["decay_ema"] is not None: + self.log_rmse_val = [] + def training_step(self, batch, batch_idx): # Get parameters ebm_opt, proposal_opt = self.optimizers() @@ -90,4 +96,9 @@ def training_step(self, batch, batch_idx): x, dic_output, ) + + if self.args_dict["dataset_name"] == "ising": + with torch.no_grad(): + log_rmse = (self.J - self.ebm.energy.J).pow(2).mean().sqrt().log() + self.log_rmse.append((log_rmse.item(), self.global_step)) return loss_total diff --git a/Model/Utils/plot_utils.py b/Model/Utils/plot_utils.py index 6324584..87c5595 100644 --- a/Model/Utils/plot_utils.py +++ b/Model/Utils/plot_utils.py @@ -7,15 +7,41 @@ import torchvision -def print_discrete_params(algo): +def print_discrete_params(algo, save_dir, step=""): if algo.args_dict["dataset_name"] == "poisson": print(f"self.ebm.energy.theta {algo.ebm.energy.lambda_}") elif algo.args_dict["dataset_name"] == "categorical": print(f"self.ebm.energy.theta {F.softmax(algo.ebm.energy.theta)}") + elif algo.args_dict["dataset_name"] == "ising": + plot_log_rmse(algo, save_dir, name="log_rmse", step=step) else: raise NotImplementedError +def plot_log_rmse(algo, save_dir, name="log_rmse", step=""): + if not os.path.exists(save_dir): + os.makedirs(save_dir) + + fig, ax = plt.subplots() + ax.plot([k[1] for k in algo.log_rmse], [k[0] for k in algo.log_rmse], "o-") + ax.set(xlabel="iteration", ylabel="log_rmse") + plt.savefig(os.path.join(save_dir, "{}_{}.png".format(name, step))) + print(f"Saved at {os.path.join(save_dir, '{}_{}.png'.format(name, step))}") + + if algo.args_dict["decay_ema"] is not None: + fig, ax = plt.subplots() + ax.plot( + [k[1] for k in algo.log_rmse_val], [k[0] for k in algo.log_rmse_val], "o-" + ) + ax.set(xlabel="iteration", ylabel="log_rmse_val") + plt.savefig( + os.path.join(save_dir, "{}_val_{}.png".format(name, step)), + ) + print( + f"Saved at {os.path.join(save_dir, '{}_val_{}.png'.format(name, step))}", + ) + + def plot_energy_2d( algo, save_dir, diff --git a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml index ff1d3c9..3b62c48 100644 --- a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml +++ b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml @@ -1,9 +1,11 @@ ebm_name : 'self_normalized' -max_epoch : 130 +max_epoch : 10 max_steps : 2000 save_energy_every : 100 -samples_every: 500 +samples_every: 2 batch_size : 64 nb_sample_bias_explicit : 1024 bias_explicit : True -num_sample_proposal : 64 \ No newline at end of file +num_sample_proposal : 512 +num_sample_proposal_val : 512 +decay_ema: 0.99 \ No newline at end of file From 6e819549acc98e015facffe075258f771ff8f4b7 Mon Sep 17 00:00:00 2001 From: Hugo Date: Mon, 15 May 2023 22:58:48 +0200 Subject: [PATCH 8/9] feat : add Ising adaptive --- .../__init__.py | 4 +- .../ising_proposal.py | 42 ++++++++++++------ .../ising_proposal_adaptive.py | 43 +++++++++++++++++++ Model/Proposals/proposal_getter.py | 2 + .../abstract_trainer.py | 1 + .../trainer_self_normalized.py | 1 + .../YAMLEBM_2D/self_normalized_ising.yaml | 6 +-- .../YAMLPROPOSAL/ising_adaptive.yaml | 4 ++ SomeBashDistributionEstimation/bash_ising.sh | 7 +++ main_trainer.py | 1 + 10 files changed, 95 insertions(+), 16 deletions(-) create mode 100644 Model/Proposals/ProposalForDistributionEstimation/ising_proposal_adaptive.py create mode 100644 Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising_adaptive.yaml create mode 100644 SomeBashDistributionEstimation/bash_ising.sh diff --git a/Model/Proposals/ProposalForDistributionEstimation/__init__.py b/Model/Proposals/ProposalForDistributionEstimation/__init__.py index 9fbf814..dd1be15 100644 --- a/Model/Proposals/ProposalForDistributionEstimation/__init__.py +++ b/Model/Proposals/ProposalForDistributionEstimation/__init__.py @@ -3,4 +3,6 @@ from .gaussian_mixture import GaussianMixtureProposal from .gaussian_mixture_adaptive import GaussianMixtureAdaptiveProposal from .noise_gradation_adaptive import NoiseGradationAdaptiveProposal -from .student import StudentProposal \ No newline at end of file +from .student import StudentProposal +from .ising_proposal import IsingProposal +from .ising_proposal_adaptive import IsingProposalAdaptive \ No newline at end of file diff --git a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py index f444a39..d0a0c30 100644 --- a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py +++ b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py @@ -13,34 +13,52 @@ class IsingProposal(nn.Module): p: float, probability of not flipping the spin """ - def __init__(self, input_size, dataset, p: float = 0.9): + def __init__(self, input_size, dataset, centers = None, p: float = 0.9): super(IsingProposal, self).__init__() self.dataset = dataset self.p = p self.dummy_param = Parameter(torch.Tensor([0.0]), requires_grad=False) + shape = self.dataset[0][0].shape + len_dataset = len(self.dataset) + if centers is not None: + # For the adaptive proposal, we can pass the centers directly + self.centers = centers + self.len_centers = len(centers) + else : + self.len_centers = len_dataset + if np.prod(shape)*len(dataset)<1e5: + # Checking the full size of storing everything : + self.centers = torch.stack([self.dataset[i][0] for i in range(len(self.dataset))]) + else : + self.centers = None + + def get_centers(self, index): + # Might be worth it time wise to store the samples in memory rather than recalculating them everytime + if self.centers is not None : + return self.centers[index] + else : + return torch.stack([self.dataset[i][0] for i in index]) + def sample(self, nb_sample: int = 1) -> Float[torch.Tensor, "nb_sample nb_nodes"]: - if nb_sample < len(self.dataset): - index = np.random.choice(len(self.dataset), nb_sample) + if nb_sample < self.len_centers: + index = np.random.choice(self.len_centers, nb_sample) else: - index = np.random.choice(len(self.dataset), nb_sample, replace=True) + index = np.random.choice(self.len_centers, nb_sample, replace=True) - center = torch.stack([self.dataset[i][0] for i in index]) + center = self.get_centers(index) bernoulli_keep = torch.distributions.Bernoulli( torch.full_like(center, self.p) ).sample() - samples = center * bernoulli_keep + (1 - center) * (1 - bernoulli_keep) + samples = center * bernoulli_keep + (1 - center) * (1 - bernoulli_keep) return samples.detach() def log_prob( self, x: Float[torch.Tensor, "batch_size nb_nodes"] ) -> Float[torch.Tensor, "batch_size"]: - data = torch.stack( - [self.dataset[i][0] for i in range(len(self.dataset))] - ).unsqueeze( - 1 - ) # len(dataset), 1, nb_nodes + + data = self.get_centers(range(self.len_centers)).unsqueeze(1) x_expanded = x.unsqueeze(0) # 1, batch_size, nb_nodes dependency = (data - x_expanded).abs() # len(dataset), batch_size, nb_nodes @@ -50,6 +68,6 @@ def log_prob( .sum(-1) ) # len(dataset), batch_size log_prob = log_prob.logsumexp(0) - torch.log( - torch.tensor(len(self.dataset)) + torch.tensor(self.len_centers) ) # batch_size return log_prob diff --git a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal_adaptive.py b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal_adaptive.py new file mode 100644 index 0000000..e04871d --- /dev/null +++ b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal_adaptive.py @@ -0,0 +1,43 @@ +import numpy as np +import torch +import torch.nn as nn +from jaxtyping import Float +from torch.nn.parameter import Parameter +from .ising_proposal import IsingProposal + +class IsingProposalAdaptive(IsingProposal): + """Proposal for Ising model + + Attributes: + dataset: torch.utils.data.Dataset, dataset of the Ising model + p: float, probability of not flipping the spin + """ + + def __init__(self, default_proposal, input_size, dataset, p: float = 0.9): + super(IsingProposalAdaptive, self).__init__(input_size =input_size, dataset=dataset, centers=None, p = p) + self.x = None + def set_x(self, x): + self.x = x + + def get_center(self, index): + # Might be worth it time wise to store the samples in memory rather than recalculating them everytime + if self.center is not None : + return self.center[index] + else : + return torch.stack([self.dataset[i][0] for i in index]) + + def sample(self, nb_sample: int = 1) -> Float[torch.Tensor, "nb_sample nb_nodes"]: + if self.x is not None : + aux_ising = IsingProposal(self.x.shape[1], self.dataset, p = self.p, centers=self.x) + return aux_ising.sample(nb_sample).detach() + else : + return super().sample(nb_sample).detach() + + def log_prob( + self, x: Float[torch.Tensor, "batch_size nb_nodes"] + ) -> Float[torch.Tensor, "batch_size"]: + if self.x is not None : + aux_ising = IsingProposal(self.x.shape[1], self.dataset, p = self.p, centers=self.x) + return aux_ising.log_prob(x) + else : + return super().log_prob(x) diff --git a/Model/Proposals/proposal_getter.py b/Model/Proposals/proposal_getter.py index f6e7aef..8b3b428 100644 --- a/Model/Proposals/proposal_getter.py +++ b/Model/Proposals/proposal_getter.py @@ -15,6 +15,7 @@ from .ProposalForDistributionEstimation.categorical import Categorical from .ProposalForDistributionEstimation.noise_gradation_adaptive import NoiseGradationAdaptiveProposal from .ProposalForDistributionEstimation.student import StudentProposal +from .ProposalForDistributionEstimation.ising_proposal_adaptive import IsingProposalAdaptive import copy dic_proposals = { @@ -25,6 +26,7 @@ "uniform_categorical": Categorical, "gaussian_mixture_adaptive": GaussianMixtureAdaptiveProposal, "ising": IsingProposal, + "ising_adaptive": IsingProposalAdaptive, 'kernel_density_adaptive': KernelDensityAdaptive, 'noise_gradation_adaptive' : NoiseGradationAdaptiveProposal, 'student' : StudentProposal, diff --git a/Model/Trainer/DistributionEstimation/abstract_trainer.py b/Model/Trainer/DistributionEstimation/abstract_trainer.py index cae62f5..9dda5ed 100644 --- a/Model/Trainer/DistributionEstimation/abstract_trainer.py +++ b/Model/Trainer/DistributionEstimation/abstract_trainer.py @@ -124,6 +124,7 @@ def validation_epoch_end(self, outputs): with torch.no_grad(): log_rmse_val = (self.J - self.ebm.energy.J).pow(2).mean().sqrt().log() self.log_rmse_val.append((log_rmse_val.item(), self.global_step)) + self.log("log_rmse_val", log_rmse_val) self.update_dic_logger(outputs, name="val_") self.proposal_visualization() self.base_dist_visualization() diff --git a/Model/Trainer/DistributionEstimation/trainer_self_normalized.py b/Model/Trainer/DistributionEstimation/trainer_self_normalized.py index 0d6a8e2..dccb89b 100644 --- a/Model/Trainer/DistributionEstimation/trainer_self_normalized.py +++ b/Model/Trainer/DistributionEstimation/trainer_self_normalized.py @@ -100,4 +100,5 @@ def training_step(self, batch, batch_idx): with torch.no_grad(): log_rmse = (self.J - self.ebm.energy.J).pow(2).mean().sqrt().log() self.log_rmse.append((log_rmse.item(), self.global_step)) + self.log("log_rmse_train", log_rmse) return loss_total diff --git a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml index 48d429f..633657d 100644 --- a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml +++ b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml @@ -6,6 +6,6 @@ samples_every: 2 batch_size : 64 nb_sample_bias_explicit : 1024 bias_explicit : True -num_sample_proposal : 512 -num_sample_proposal_val : 512 -decay_ema: 0.99 +num_sample_proposal : 1024 +num_sample_proposal_val : 2048 +decay_ema: 0.9 diff --git a/Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising_adaptive.yaml b/Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising_adaptive.yaml new file mode 100644 index 0000000..5f7f110 --- /dev/null +++ b/Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising_adaptive.yaml @@ -0,0 +1,4 @@ +proposal_name: "ising_adaptive" +default_proposal_name : "ising" +proposal_params: + p: 0.5 diff --git a/SomeBashDistributionEstimation/bash_ising.sh b/SomeBashDistributionEstimation/bash_ising.sh new file mode 100644 index 0000000..e881ab8 --- /dev/null +++ b/SomeBashDistributionEstimation/bash_ising.sh @@ -0,0 +1,7 @@ +python main_trainer.py \ +--yamldataset Dataset/MissingDataDataset/YAMLExamples/ising.yaml \ +--yamlebm Model/YAMLDISTRIBUTION/YAMLBASEDIST/none.yaml \ +Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml \ +Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising.yaml \ +Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml \ +Model/YAMLDISTRIBUTION/YAMLOPTIMIZATION/adam1e-3.yaml \ No newline at end of file diff --git a/main_trainer.py b/main_trainer.py index 55f66c3..0e93e36 100644 --- a/main_trainer.py +++ b/main_trainer.py @@ -165,6 +165,7 @@ def find_last_version(dir): if not args_dict['just_test']: + trainer.validate(algo, dataloaders=val_loader) trainer.fit(algo, train_dataloaders=train_loader, val_dataloaders=val_loader) algo.load_state_dict( torch.load(checkpoint_callback_val.best_model_path)["state_dict"] From a142239d6aefc70be90ef1f83e84f287bd7288be Mon Sep 17 00:00:00 2001 From: HugoSenetaire Date: Mon, 15 May 2023 23:15:11 +0200 Subject: [PATCH 9/9] feat : use large proposal on gpu --- Dataset/MissingDataDataset | 2 +- .../ising_proposal.py | 29 ++++++++++--------- .../abstract_trainer.py | 2 +- .../trainer_self_normalized.py | 6 ++-- .../YAMLEBM_2D/self_normalized_ising.yaml | 4 +-- ...self_normalized_ising_verylargesample.yaml | 11 +++++++ SomeBashDistributionEstimation/bash_ising.sh | 10 ++++++- 7 files changed, 42 insertions(+), 22 deletions(-) create mode 100644 Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising_verylargesample.yaml diff --git a/Dataset/MissingDataDataset b/Dataset/MissingDataDataset index 4791d39..dc3ab09 160000 --- a/Dataset/MissingDataDataset +++ b/Dataset/MissingDataDataset @@ -1 +1 @@ -Subproject commit 4791d397ce12ef0d90cf07493b1d436da2b12607 +Subproject commit dc3ab097ba4ab09749afcb60b2c6fbbd311c1f49 diff --git a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py index d0a0c30..aa82dd4 100644 --- a/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py +++ b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py @@ -27,11 +27,11 @@ def __init__(self, input_size, dataset, centers = None, p: float = 0.9): self.len_centers = len(centers) else : self.len_centers = len_dataset - if np.prod(shape)*len(dataset)<1e5: + # if np.prod(shape)*len(dataset)<1e5: # Checking the full size of storing everything : - self.centers = torch.stack([self.dataset[i][0] for i in range(len(self.dataset))]) - else : - self.centers = None + # self.centers = torch.stack([self.dataset[i][0] for i in range(len(self.dataset))]) + # else : + self.centers = None def get_centers(self, index): # Might be worth it time wise to store the samples in memory rather than recalculating them everytime @@ -41,24 +41,25 @@ def get_centers(self, index): return torch.stack([self.dataset[i][0] for i in index]) def sample(self, nb_sample: int = 1) -> Float[torch.Tensor, "nb_sample nb_nodes"]: - if nb_sample < self.len_centers: - index = np.random.choice(self.len_centers, nb_sample) - else: - index = np.random.choice(self.len_centers, nb_sample, replace=True) + with torch.no_grad(): # Lower memory size + if nb_sample < self.len_centers: + index = np.random.choice(self.len_centers, nb_sample) + else: + index = np.random.choice(self.len_centers, nb_sample, replace=True) - center = self.get_centers(index) - bernoulli_keep = torch.distributions.Bernoulli( - torch.full_like(center, self.p) - ).sample() + center = self.get_centers(index).to(self.dummy_param.device) + bernoulli_keep = torch.distributions.Bernoulli( + torch.full_like(center, self.p) + ).sample() - samples = center * bernoulli_keep + (1 - center) * (1 - bernoulli_keep) + samples = center * bernoulli_keep + (1 - center) * (1 - bernoulli_keep) return samples.detach() def log_prob( self, x: Float[torch.Tensor, "batch_size nb_nodes"] ) -> Float[torch.Tensor, "batch_size"]: - data = self.get_centers(range(self.len_centers)).unsqueeze(1) + data = self.get_centers(range(self.len_centers)).unsqueeze(1).to(x.device) x_expanded = x.unsqueeze(0) # 1, batch_size, nb_nodes dependency = (data - x_expanded).abs() # len(dataset), batch_size, nb_nodes diff --git a/Model/Trainer/DistributionEstimation/abstract_trainer.py b/Model/Trainer/DistributionEstimation/abstract_trainer.py index 9dda5ed..1a5c26b 100644 --- a/Model/Trainer/DistributionEstimation/abstract_trainer.py +++ b/Model/Trainer/DistributionEstimation/abstract_trainer.py @@ -122,7 +122,7 @@ def validation_step(self, batch, batch_idx): def validation_epoch_end(self, outputs): if self.args_dict["dataset_name"] == "ising": with torch.no_grad(): - log_rmse_val = (self.J - self.ebm.energy.J).pow(2).mean().sqrt().log() + log_rmse_val = (self.J.to(self.device, self.dtype) - self.ebm.energy.J).pow(2).mean().sqrt().log() self.log_rmse_val.append((log_rmse_val.item(), self.global_step)) self.log("log_rmse_val", log_rmse_val) self.update_dic_logger(outputs, name="val_") diff --git a/Model/Trainer/DistributionEstimation/trainer_self_normalized.py b/Model/Trainer/DistributionEstimation/trainer_self_normalized.py index dccb89b..fb3523d 100644 --- a/Model/Trainer/DistributionEstimation/trainer_self_normalized.py +++ b/Model/Trainer/DistributionEstimation/trainer_self_normalized.py @@ -35,7 +35,7 @@ def __init__( ) if args_dict["dataset_name"] == "ising": - self.J = complete_dataset.J + self.J = complete_dataset.J.to(self.device) self.log_rmse = [] if args_dict["decay_ema"] is not None: self.log_rmse_val = [] @@ -98,7 +98,7 @@ def training_step(self, batch, batch_idx): if self.args_dict["dataset_name"] == "ising": with torch.no_grad(): - log_rmse = (self.J - self.ebm.energy.J).pow(2).mean().sqrt().log() - self.log_rmse.append((log_rmse.item(), self.global_step)) + log_rmse = (self.J.to(self.device, self.dtype) - self.ebm.energy.J).pow(2).mean().sqrt().log() + self.log_rmse.append((log_rmse.detach().item(), self.global_step)) self.log("log_rmse_train", log_rmse) return loss_total diff --git a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml index 633657d..a8017e9 100644 --- a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml +++ b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml @@ -1,5 +1,5 @@ ebm_name : 'self_normalized' -max_epoch : 10 +max_epoch : 100 max_steps : 2000 save_energy_every : 100 samples_every: 2 @@ -7,5 +7,5 @@ batch_size : 64 nb_sample_bias_explicit : 1024 bias_explicit : True num_sample_proposal : 1024 -num_sample_proposal_val : 2048 +num_sample_proposal_val : 1024 decay_ema: 0.9 diff --git a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising_verylargesample.yaml b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising_verylargesample.yaml new file mode 100644 index 0000000..02cbdb3 --- /dev/null +++ b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising_verylargesample.yaml @@ -0,0 +1,11 @@ +ebm_name : 'self_normalized' +max_epoch : 100 +max_steps : 2000 +save_energy_every : 100 +samples_every: 2 +batch_size : 64 +nb_sample_bias_explicit : 1024 +bias_explicit : True +num_sample_proposal : 2048 +num_sample_proposal_val : 2048 +decay_ema: 0.9 diff --git a/SomeBashDistributionEstimation/bash_ising.sh b/SomeBashDistributionEstimation/bash_ising.sh index e881ab8..02c6548 100644 --- a/SomeBashDistributionEstimation/bash_ising.sh +++ b/SomeBashDistributionEstimation/bash_ising.sh @@ -1,7 +1,15 @@ -python main_trainer.py \ +CUDA_VISIBLE_DEVICES=1 python main_trainer.py \ --yamldataset Dataset/MissingDataDataset/YAMLExamples/ising.yaml \ --yamlebm Model/YAMLDISTRIBUTION/YAMLBASEDIST/none.yaml \ Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml \ Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising.yaml \ Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml \ +Model/YAMLDISTRIBUTION/YAMLOPTIMIZATION/adam1e-3.yaml + +CUDA_VISIBLE_DEVICES=2 python main_trainer.py \ +--yamldataset Dataset/MissingDataDataset/YAMLExamples/ising.yaml \ +--yamlebm Model/YAMLDISTRIBUTION/YAMLBASEDIST/none.yaml \ +Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising_verylargesample.yaml \ +Model/YAMLDISTRIBUTION/YAMLPROPOSAL/ising.yaml \ +Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml \ Model/YAMLDISTRIBUTION/YAMLOPTIMIZATION/adam1e-3.yaml \ No newline at end of file