layers.py 🚧¶
Source: LION/models/diffusion/NCSNpp_helpers/layers.py
Warning
This file has not yet received a complete narrative and docstring audit. Its public source-level API is listed automatically below.
Common layers for defining score networks.
- get_act(params)
Get activation functions from the config file.
- ncsn_conv1x1(in_planes, out_planes, stride=1, bias=True, dilation=1, init_scale=1.0, padding=0)
1x1 convolution. Same as NCSNv1/v2.
- variance_scaling(scale, mode, distribution, in_axis=1, out_axis=0, dtype=torch.float32, device='cpu')
Ported from JAX.
- default_init(scale=1.0)
The same initialization used in DDPM.
- class Dense
Linear layer with default_init.
- ddpm_conv1x1(in_planes, out_planes, stride=1, bias=True, init_scale=1.0, padding=0)
1x1 convolution with DDPM initialization.
- ncsn_conv3x3(in_planes, out_planes, stride=1, bias=True, dilation=1, init_scale=1.0, padding=1)
3x3 convolution with PyTorch initialization. Same as NCSNv1/NCSNv2.
- ddpm_conv3x3(in_planes, out_planes, stride=1, bias=True, dilation=1, init_scale=1.0, padding=1)
3x3 convolution with DDPM initialization.
- class CRPBlock(features, n_stages, act=nn.ReLU(), maxpool=True)
No docstring is available.
- CRPBlock.forward(self, x)
No docstring is available.
- class CondCRPBlock(features, n_stages, num_classes, normalizer, act=nn.ReLU())
No docstring is available.
- CondCRPBlock.forward(self, x, y)
No docstring is available.
- class RCUBlock(features, n_blocks, n_stages, act=nn.ReLU())
No docstring is available.
- RCUBlock.forward(self, x)
No docstring is available.
- class CondRCUBlock(features, n_blocks, n_stages, num_classes, normalizer, act=nn.ReLU())
No docstring is available.
- CondRCUBlock.forward(self, x, y)
No docstring is available.
- class MSFBlock(in_planes, features)
No docstring is available.
- MSFBlock.forward(self, xs, shape)
No docstring is available.
- class CondMSFBlock(in_planes, features, num_classes, normalizer)
No docstring is available.
- CondMSFBlock.forward(self, xs, y, shape)
No docstring is available.
- class RefineBlock(in_planes, features, act=nn.ReLU(), start=False, end=False, maxpool=True)
No docstring is available.
- RefineBlock.forward(self, xs, output_shape)
No docstring is available.
- class CondRefineBlock(in_planes, features, num_classes, normalizer, act=nn.ReLU(), start=False, end=False)
No docstring is available.
- CondRefineBlock.forward(self, xs, y, output_shape)
No docstring is available.
- class ConvMeanPool(input_dim, output_dim, kernel_size=3, biases=True, adjust_padding=False)
No docstring is available.
- ConvMeanPool.forward(self, inputs)
No docstring is available.
- class MeanPoolConv(input_dim, output_dim, kernel_size=3, biases=True)
No docstring is available.
- MeanPoolConv.forward(self, inputs)
No docstring is available.
- class UpsampleConv(input_dim, output_dim, kernel_size=3, biases=True)
No docstring is available.
- UpsampleConv.forward(self, inputs)
No docstring is available.
- class ConditionalResidualBlock(input_dim, output_dim, num_classes, resample=1, act=nn.ELU(), normalization=ConditionalInstanceNorm2dPlus, adjust_padding=False, dilation=None)
No docstring is available.
- ConditionalResidualBlock.forward(self, x, y)
No docstring is available.
- class ResidualBlock(input_dim, output_dim, resample=None, act=nn.ELU(), normalization=nn.InstanceNorm2d, adjust_padding=False, dilation=1)
No docstring is available.
- ResidualBlock.forward(self, x)
No docstring is available.
- get_timestep_embedding(timesteps, embedding_dim, max_positions=10000)
No docstring is available.
- contract_inner(x, y)
tensordot(x, y, 1).
- class NIN(in_dim, num_units, init_scale=0.1)
No docstring is available.
- NIN.forward(self, x)
No docstring is available.
- class AttnBlock(channels)
Channel-wise self-attention block.
- AttnBlock.forward(self, x)
No docstring is available.
- class Upsample(channels, with_conv=False)
No docstring is available.
- Upsample.forward(self, x)
No docstring is available.
- class Downsample(channels, with_conv=False)
No docstring is available.
- Downsample.forward(self, x)
No docstring is available.
- class ResnetBlockDDPM(act, in_ch, out_ch=None, temb_dim=None, conv_shortcut=False, dropout=0.1)
The ResNet Blocks used in DDPM.
- ResnetBlockDDPM.forward(self, x, temb=None)
No docstring is available.