core.layer

Overview

Contains probabilistic versions of common neural network layers and the infrastructure to attach distributions to them.

Contents

  • AbstractProbLayer: A base class managing probabilistic_mode and sampling
  • ProbConv2d, ProbLinear, ProbBatchNorm1d, ProbBatchNorm2d: Probabilistic layers derived from PyTorch
  • utils.py for layer inspection and traversal

These layers replace deterministic parameters with sampled ones, supporting Bayesian inference under the PAC-Bayes framework.

 1"""
 2## Overview
 3Contains probabilistic versions of common neural network layers and the
 4infrastructure to attach distributions to them.
 5
 6## Contents
 7- **AbstractProbLayer**: A base class managing probabilistic_mode and sampling
 8- **ProbConv2d, ProbLinear, ProbBatchNorm1d, ProbBatchNorm2d**:
 9  Probabilistic layers derived from PyTorch
10- **utils.py** for layer inspection and traversal
11
12These layers replace deterministic parameters with sampled ones,
13supporting Bayesian inference under the PAC-Bayes framework.
14"""
15
16from core.layer.AbstractProbLayer import AbstractProbLayer
17from core.layer.ProbBatchNorm1d import ProbBatchNorm1d
18from core.layer.ProbBatchNorm2d import ProbBatchNorm2d
19from core.layer.ProbConv2d import ProbConv2d
20from core.layer.ProbLinear import ProbLinear
21from core.layer.supported_layers import LAYER_MAPPING