Including minibatch MCMC for paper 2

Author: papamarkou

.gitignore

@@ -9,6 +9,7 @@ eeyore/data/__pycache__/
 eeyore/datasets/__pycache__/
 eeyore/distributions/__pycache__/
 eeyore/integrators/__pycache__/
+eeyore/itertools/__pycache__/
 eeyore/kernels/__pycache__/
 eeyore/linalg/__pycache__/
 eeyore/models/__pycache__/

eeyore/chains/chain_file.py

@@ -25,16 +25,19 @@ def close(self):
         for key in self.vals.keys():
             self.vals[key].close()
 
-    def update(self, state, reset=True, close=True):
+    def update(self, state,
+        reset=True, close=True, fmt={'sample': '%.18e', 'target_val': '%.18e', 'grad_val': '%.18e', 'accepted': '%d'}):
         """ Update the chain """
         if reset:
             self.reset(keys=self.vals.keys())
 
         for key in self.vals.keys():
             if isinstance(state[key], torch.Tensor):
-                np.savetxt(self.vals[key], state[key].detach().cpu().numpy().ravel()[np.newaxis], delimiter=',')
+                np.savetxt(
+                    self.vals[key], state[key].detach().cpu().numpy().ravel()[np.newaxis], fmt=fmt[key], delimiter=','
+                )
             elif isinstance(state[key], np.ndarray):
-                np.savetxt(self.vals[key], state[key].ravel()[np.newaxis], delimiter=',')
+                np.savetxt(self.vals[key], state[key].ravel()[np.newaxis], fmt=fmt[key], delimiter=',')
             else:
                 self.vals[key].write(str(state[key])+'\n')
 

eeyore/chains/chain_list.py

@@ -92,9 +92,12 @@ def mc_cor(self, mc_cov_mat=None, method='inse', adjust=False):
             return st.cor_from_cov(mc_cov_mat)
 
     def acceptance_rate(self):
-        """ proportion of accepted samples """
+        """ Proportion of accepted samples """
         return sum(self.vals['accepted']) / self.num_samples()
 
+    def block_acceptance_rate(self):
+        return torch.stack(self.vals['accepted']).sum(axis=0) / self.num_samples()
+
     def multi_ess(self, mc_cov_mat=None, method='inse', adjust=False):
         return st.multi_ess(self.get_samples(), mc_cov_mat=mc_cov_mat, method=method, adjust=adjust)
 
@@ -106,13 +109,17 @@ def load(self, path):
         """ Load a previously saved chain """
         self.vals = torch.load(path)
 
-    def to_chainfile(self, keys=None, path=Path.cwd(), mode='a'):
+    def to_chainfile(self,
+        keys=None,
+        path=Path.cwd(),
+        mode='a',
+        fmt={'sample': '%.18e', 'target_val': '%.18e', 'grad_val': '%.18e', 'accepted': '%d'}):
         from .chain_file import ChainFile
 
         chainfile = ChainFile(keys=keys or self.vals.keys(), path=path, mode=mode)
 
         for i in range(len(self)):
-            chainfile.update(self.state(i), reset=False, close=False)
+            chainfile.update(self.state(i), reset=False, close=False, fmt=fmt)
 
         chainfile.close()
 

eeyore/datasets/__init__.py

@@ -2,5 +2,7 @@
 from .data_info import data_paths
 from .empty_dataset import EmptyXYDataset
 from .idataset import IDataset
+from .mld_batcher import MLDBatcher
+from .mld_classification_batcher import MLDClassificationBatcher
 from .xydataset import XYDataset
 from .xyidataset import XYIDataset

eeyore/datasets/mld_batcher.py

@@ -0,0 +1,4 @@
+# Minimum likelihood distance batcher (abstract base class)
+
+class MLDBatcher:
+    pass

eeyore/datasets/mld_classification_batcher.py

@@ -0,0 +1,90 @@
+# Minimum likelihood distance batcher for classification
+
+import copy
+import functools
+import operator
+import random
+import torch
+
+from .mld_batcher import MLDBatcher
+
+class MLDClassificationBatcher(MLDBatcher):
+    def __init__(self, num_batches, chunk_sizes, dataset=None):
+        self.num_batches = num_batches
+
+        self.chunk_sizes = chunk_sizes
+        assert len(self.chunk_sizes) == 2
+
+        self.set_dataset(dataset)
+
+    def set_dataset(self, dataset):
+        self.dataset = dataset
+
+        if self.dataset is not None:
+            self.num_points = len(dataset)
+            self.num_classes = len(dataset.y[0])
+
+            label_argmax = torch.argmax(self.dataset.y, axis=1)
+
+            self.class_indices = [[] for _ in range(self.num_classes)]
+            for i in range(self.num_points):
+                self.class_indices[label_argmax[i].item()].append(i)
+
+            self.class_props = [len(self.class_indices[i]) / self.num_points for i in range(self.num_classes)]
+
+            self.class_num_batch_points = [
+                [int(self.class_props[j]*self.chunk_sizes[i]) for j in range(self.num_classes)] for i in range(2)
+            ]
+
+    def batch_size(self):
+        return sum(self.chunk_sizes)
+
+    def fill_class_sizes(self):
+        class_num_batch_points = copy.deepcopy(self.class_num_batch_points)
+
+        sampled_classes = [
+            random.choices(range(self.num_classes), k=self.chunk_sizes[i]-sum(class_num_batch_points[i])) for i in range(2)
+        ]
+
+        for i in range(2):
+            for j in sampled_classes[i]:
+                class_num_batch_points[i][j] = class_num_batch_points[i][j] + 1
+
+        return class_num_batch_points
+
+    def get_batch(self, model, params, fill=True):
+        class_num_batch_points = [self.fill_class_sizes() for _ in range(self.num_batches)]
+
+        mld_distance = float('inf')
+
+        for i in range(self.num_batches):
+            indices = []
+
+            indices.extend([
+                random.sample(self.class_indices[j], class_num_batch_points[i][0][j]) for j in range(self.num_classes)
+            ])
+
+            rest_indices = [list(set(self.class_indices[j]) - set(indices[j])) for j in range(self.num_classes)]
+
+            indices.extend(
+                [random.sample(rest_indices[j], class_num_batch_points[i][1][j]) for j in range(self.num_classes)]
+            )
+
+            indices = functools.reduce(operator.iconcat, indices, [])
+
+            indices.sort()
+
+            distance = 0.
+
+            for j in range(2):
+                log_lik_vals = model.set_params_and_log_lik(params[j].clone().detach(), self.dataset.x, self.dataset.y)
+
+                distance = distance + (log_lik_vals.mean() - log_lik_vals[indices].mean()).abs()
+
+            distance = distance.sqrt().item()
+
+            if distance < mld_distance:
+                mld_indices = indices
+                mld_distance = distance
+
+        return self.dataset.x[mld_indices, :], self.dataset.y[mld_indices, :]

eeyore/itertools/__init__.py

@@ -0,0 +1 @@
+from .chunk_evenly import chunk_evenly

eeyore/itertools/chunk_evenly.py

@@ -0,0 +1,9 @@
+def chunk_evenly(iterable, n):
+    iterable_len = len(iterable)
+    r, a = iterable_len % n, 0
+    
+    for i, s in enumerate(range(0, iterable_len if (r == 0) else (iterable_len-n), n)):
+        yield iterable[(s+a):((s+a+n+1) if (i < r) else (s+a+n))]
+
+        if i < r:
+            a = a + 1

eeyore/models/mlp.py

@@ -1,3 +1,4 @@
+import itertools
 import torch
 import torch.nn as nn
 
@@ -51,3 +52,47 @@ def forward(self, x):
     def num_hidden_layers(self):
         """ Get the number of hidden layers. """
         return len(self.hp.dims)-2
+
+    def num_par_blocks(self):
+        return sum(self.hp.dims[1:])
+
+    def layer_and_node_from_par_block(self, b):
+        num_nodes_per_layer = [0] + list(itertools.accumulate(self.hp.dims[1:]))
+        l = self.num_hidden_layers()
+
+        for i in range(1, len(num_nodes_per_layer)):
+            if num_nodes_per_layer[-i-1] <= b < num_nodes_per_layer[-i]:
+                n = b if (num_nodes_per_layer[-i-1] == 0) else (b % num_nodes_per_layer[-i-1])
+                break
+            else:
+                l = l - 1
+
+        return l, n
+
+    def starting_par_block_idx(self, l):
+        s = 0
+
+        if l > 0:
+            for i in range(l):
+                s = s + (self.hp.dims[i]+1 if self.hp.bias[i] else self.hp.dims[i])*self.hp.dims[i+1]
+
+        return s
+
+    def starting_par_block_indices(self):
+        s = [0]
+
+        for l in range(self.num_hidden_layers()):
+            s.append(s[-1]+(self.hp.dims[l]+1 if self.hp.bias[l] else self.hp.dims[l])*self.hp.dims[l+1])
+
+        return s
+
+    def par_block_indices(self, b):
+        l, n = self.layer_and_node_from_par_block(b)
+        s = self.starting_par_block_idx(l)
+
+        indices = list(range(s+n*self.hp.dims[l], s+(n+1)*self.hp.dims[l])) if (self.hp.dims[l] > 1) else [s+n]
+
+        if self.hp.bias[l]:
+            indices.append(s+self.hp.dims[l]*self.hp.dims[l+1]+n)
+
+        return indices, l, n

eeyore/models/tmp.py

@@ -0,0 +1,27 @@
+# %%
+
+from eeyore.constants import loss_functions
+from eeyore.models import mlp
+
+# %%
+
+# hparams = mlp.Hyperparameters(dims=[2, 3, 3, 2], bias=3*[True], activations=3*[None])
+hparams = mlp.Hyperparameters(dims=[2, 3, 3, 2], bias=[True, True, True], activations=3*[None])
+
+model = mlp.MLP(loss=loss_functions['multiclass_classification'], hparams=hparams)
+
+# %%
+
+print(model.num_par_blocks())
+
+print([model.starting_par_block_idx(i) for i in [0, 1, 2]])
+
+print(model.starting_par_block_indices())
+
+for b in range(8):
+    l, n = model.layer_and_node_from_par_block(b)
+    print("Block {} is in layer {} and node {} of that layer".format(b, l, n))
+
+for b in range(8):
+    indices, l, n = model.par_block_indices(b)
+    print("Block {} is in layer {} and node {} of that layer and has indices {}".format(b, l, n, indices))

eeyore/samplers/__init__.py

@@ -1,4 +1,6 @@
 from .am import AM
+from .dmcl import DMCL
+from .gibbs import Gibbs
 from .hmc import HMC
 from .mala import MALA
 from .metropolis_hastings import MetropolisHastings

eeyore/samplers/dmcl.py

@@ -0,0 +1,109 @@
+import json
+import torch
+
+from .single_chain_serial_sampler import SingleChainSerialSampler
+from eeyore.chains import ChainList
+from eeyore.datasets import DataCounter
+from eeyore.itertools import chunk_evenly
+from eeyore.kernels import NormalKernel
+
+class DMCL(SingleChainSerialSampler):
+    def __init__(self, model, batch_model, batcher,
+        theta0=None, dataloader=None, data0=None, counter=None,
+        scales=1., node_subblock_size=None, chain=ChainList()):
+        super(DMCL, self).__init__(counter or DataCounter.from_dataloader(dataloader))
+        self.model = model
+        self.batch_model = batch_model
+        self.dataloader = dataloader
+
+        self.batcher = batcher
+        self.batcher.set_dataset(dataloader.dataset)
+
+        self.keys = ['sample', 'target_val', 'accepted']
+        self.chain = chain
+
+        if theta0 is not None:
+            self.set_current(theta0.clone().detach(), data=data0)
+
+        if isinstance(scales, float):
+            self.scales = torch.full([self.model.num_par_blocks()], scales, dtype=self.model.dtype, device=self.model.device)
+        elif isinstance(scales, torch.Tensor):
+            self.scales = scales.to(dtype=self.model.dtype, device=self.model.device)
+        elif isinstance(scales, list):
+            self.scales = torch.tensor(scales, dtype=self.model.dtype, device=self.model.device)
+        else:
+            self.scales = scales
+
+        if node_subblock_size is None:
+            self.node_subblock_size = [None for _ in range(self.model.num_par_blocks())]
+        else:
+            self.node_subblock_size = node_subblock_size
+
+    def set_current(self, theta, data=None):
+        super().set_current(theta, data=data)
+
+    def reset(self, theta, data=None, reset_counter=True, reset_chain=True):
+        super().reset(theta, data=data, reset_counter=reset_counter, reset_chain=reset_chain)
+
+    def get_blocks(self):
+        blocks = []
+
+        for b in range(self.model.num_par_blocks()):
+            indices, l, n = self.model.par_block_indices(b)
+
+            if self.node_subblock_size[b] is None:
+                indices = [indices]
+            else:
+                indices = list(chunk_evenly(indices, self.node_subblock_size[b]))
+
+            blocks.append([l, n, indices])
+
+        return blocks
+
+    def save_blocks(self, path='gibbs_lbocks.txt', mode='w'):
+        with open(path, mode) as file:
+            json.dump(self.get_blocks(), file)
+
+    def draw(self, x, y, savestate=False):
+        proposed = {key : None for key in self.keys}
+        self.current['accepted'] = []
+
+        proposed['sample'] = self.current['sample'].clone().detach()
+
+        for b in range(self.model.num_par_blocks()):
+            indices, _, _ = self.model.par_block_indices(b)
+
+            if self.node_subblock_size[b] is None:
+                indices = [indices]
+            else:
+                indices = list(chunk_evenly(indices, self.node_subblock_size[b]))
+
+            for i in range(len(indices)):
+                kernel = NormalKernel(proposed['sample'][indices[i]], self.scales[b])
+
+                proposed['sample'][indices[i]] = kernel.sample()
+
+                x_batch, y_batch = self.batcher.get_batch(
+                    self.batch_model,
+                    [self.current['sample'].clone().detach(), proposed['sample'].clone().detach()],
+                    fill=True
+                )
+
+                self.current['target_val'] = self.model.log_target(self.current['sample'].clone().detach(), x_batch, y_batch)
+                proposed['target_val'] = self.model.log_target(proposed['sample'].clone().detach(), x_batch, y_batch)
+
+                log_rate = proposed['target_val'] - self.current['target_val']
+                if torch.log(torch.rand(1, dtype=self.model.dtype, device=self.model.device)) < log_rate:
+                    self.current['sample'][indices[i]] = proposed['sample'][indices[i]]
+                    self.current['accepted'].append(1)
+                else:
+                    self.model.set_params(self.current['sample'].clone().detach())
+                    self.current['accepted'].append(0)
+
+        self.current['accepted'] = torch.tensor(self.current['accepted'], device=self.model.device)
+
+        if savestate:
+            self.chain.detach_and_update(self.current)
+
+        self.current['sample'].detach_()
+        self.current['target_val'].detach_()