docs: E-MCMC

kozistr · May 5, 2024 · 8c8b821 · 8c8b821
1 parent 2da8916
commit 8c8b821
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/docs/util.md b/docs/util.md
@@ -84,6 +84,10 @@
     :docstring:
     :members:
 
+::: pytorch_optimizer.optimizer.utils.reg_noise
+    :docstring:
+    :members:
+
 ## Newton methods
 
 ::: pytorch_optimizer.optimizer.shampoo_utils.power_iteration

diff --git a/pytorch_optimizer/optimizer/utils.py b/pytorch_optimizer/optimizer/utils.py
@@ -283,6 +283,17 @@ def reduce_max_except_dim(x: torch.Tensor, dim: int) -> torch.Tensor:
 def reg_noise(
     network1: nn.Module, network2: nn.Module, num_data: int, lr: float, eta: float = 8e-3, temperature: float = 1e-4
 ) -> torch.Tensor | float:
+    r"""Entropy-MCMC: Sampling from flat basins with ease.
+
+    usage: https://github.com/lblaoke/EMCMC/blob/master/exp/cifar10_emcmc.py
+
+    :param network1: nn.Module. network.
+    :param network2: nn.Module. network.
+    :param num_data: int. number of training data.
+    :param lr: float. learning rate.
+    :param eta: float. eta.
+    :param temperature: float. temperature.
+    """
     reg_coef: float = 0.5 / (eta * num_data)
     noise_coef: float = math.sqrt(2.0 / lr / num_data * temperature)