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 diff --git a/Dataset/MissingDataDataset b/Dataset/MissingDataDataset index 9c3a5bd..dc3ab09 160000 --- a/Dataset/MissingDataDataset +++ b/Dataset/MissingDataDataset @@ -1 +1 @@ -Subproject commit 9c3a5bd6ba0ac8dfd8cab9dc1d53ff3b6dc42e0f +Subproject commit dc3ab097ba4ab09749afcb60b2c6fbbd311c1f49 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 4a114bb..e2a3631 100644 --- a/Model/Energy/EnergyForDistribution/ising.py +++ b/Model/Energy/EnergyForDistribution/ising.py @@ -1,13 +1,16 @@ 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 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 +28,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. @@ -34,28 +38,48 @@ class EnergyIsing(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__() - 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 + 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() + 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() + + 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 + + @property + def J(self): + return self.G 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.J + xgx = (xg * x).sum(-1) + b = (self.bias[None, :] * x).sum(-1) + return -xgx - b 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/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/__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/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 new file mode 100644 index 0000000..aa82dd4 --- /dev/null +++ b/Model/Proposals/ProposalForDistributionEstimation/ising_proposal.py @@ -0,0 +1,74 @@ +import numpy as np +import torch +import torch.nn as nn +from jaxtyping import Float +from torch.nn.parameter import Parameter + + +class IsingProposal(nn.Module): + """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, 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"]: + 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).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) + 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).to(x.device) + 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(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 01f1680..8b3b428 100644 --- a/Model/Proposals/proposal_getter.py +++ b/Model/Proposals/proposal_getter.py @@ -1,16 +1,21 @@ -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 from .ProposalForDistributionEstimation.noise_gradation_adaptive import NoiseGradationAdaptiveProposal from .ProposalForDistributionEstimation.student import StudentProposal +from .ProposalForDistributionEstimation.ising_proposal_adaptive import IsingProposalAdaptive import copy dic_proposals = { @@ -19,14 +24,15 @@ "gaussian_mixture": GaussianMixtureProposal, "poisson": Poisson, "uniform_categorical": Categorical, + "gaussian_mixture_adaptive": GaussianMixtureAdaptiveProposal, + "ising": IsingProposal, + "ising_adaptive": IsingProposalAdaptive, 'kernel_density_adaptive': KernelDensityAdaptive, - 'gaussian_mixture_adaptive': GaussianMixtureAdaptiveProposal, 'noise_gradation_adaptive' : NoiseGradationAdaptiveProposal, 'student' : StudentProposal, } - def get_proposal(args_dict, input_size, dataset): if isinstance(dataset, list): @@ -34,12 +40,16 @@ def get_proposal(args_dict, input_size, dataset): else : current_dataset = 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'] = {} @@ -55,19 +65,27 @@ def get_proposal(args_dict, input_size, dataset): 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/Trainer/DistributionEstimation/abstract_trainer.py b/Model/Trainer/DistributionEstimation/abstract_trainer.py index 751993d..1a5c26b 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 self.num_samples_val = args_dict["num_sample_proposal_val"] @@ -117,9 +116,15 @@ 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.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_") self.proposal_visualization() self.base_dist_visualization() @@ -475,30 +480,63 @@ 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 self.input_type == 'image': - 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 + if self.args_dict["dataset_name"] in dic_discrete_dataset.keys(): + print_discrete_params(self, save_dir=save_dir, step=self.global_step) + else : + 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 self.input_type == 'image': + 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: + samples, init_samples = self.samples_mcmc(num_samples=num_samples) + self._sample_categorical( + num_samples=num_samples, + save_dir=save_dir, + name="samples", + step=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 ae01d14..fb3523d 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). @@ -33,6 +34,12 @@ def __init__( **kwargs, ) + if args_dict["dataset_name"] == "ising": + self.J = complete_dataset.J.to(self.device) + 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_perso() @@ -42,8 +49,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,4 +96,9 @@ def training_step(self, batch, batch_idx): dic_output, ) + if self.args_dict["dataset_name"] == "ising": + with torch.no_grad(): + 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/Utils/plot_utils.py b/Model/Utils/plot_utils.py index 53ff74b..cdabeea 100644 --- a/Model/Utils/plot_utils.py +++ b/Model/Utils/plot_utils.py @@ -3,9 +3,43 @@ 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, 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( diff --git a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml index 0573bc2..c6127be 100644 --- a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml +++ b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized.yaml @@ -5,4 +5,4 @@ samples_every: 500 batch_size : 512 nb_sample_bias_explicit : 1024 bias_explicit : True -num_sample_proposal : 512 \ No newline at end of file +num_sample_proposal : 512 diff --git a/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.yaml new file mode 100644 index 0000000..a8017e9 --- /dev/null +++ b/Model/YAMLDISTRIBUTION/YAMLEBM_2D/self_normalized_ising.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 : 1024 +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/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml b/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml index 8b077d5..fda5b2b 100644 --- a/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml +++ b/Model/YAMLDISTRIBUTION/YAMLENERGY/energy_ising.yaml @@ -1,4 +1,7 @@ energy_name: 'ising' -energy_params : - 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 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 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..02c6548 --- /dev/null +++ b/SomeBashDistributionEstimation/bash_ising.sh @@ -0,0 +1,15 @@ +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 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"]