ding.torch_utils.network.nn_module

ding.torch_utils.network.nn_module

ChannelShuffle

Bases: Module

Overview

Apply channel shuffle to the input tensor. For more details about the channel shuffle, please refer to the 'ShuffleNet' paper: https://arxiv.org/abs/1707.01083

Interfaces: __init__, forward

__init__(group_num)

Overview

Initialize the ChannelShuffle class.

Arguments: - group_num (:obj:int): The number of groups to exchange.

forward(x)

Overview

Forward pass through the ChannelShuffle module.

Arguments: - x (:obj:torch.Tensor): The input tensor. Returns: - x (:obj:torch.Tensor): The shuffled input tensor.

NearestUpsample

Bases: Module

Overview

This module upsamples the input to the given scale_factor using the nearest mode.

Interfaces: __init__, forward

__init__(scale_factor)

Overview

Initialize the NearestUpsample class.

Arguments: - scale_factor (:obj:Union[float, List[float]]): The multiplier for the spatial size.

forward(x)

Overview

Return the upsampled input tensor.

Arguments: - x (:obj:torch.Tensor): The input tensor. Returns: - upsample(:obj:torch.Tensor): The upsampled input tensor.

BilinearUpsample

Bases: Module

Overview

This module upsamples the input to the given scale_factor using the bilinear mode.

Interfaces: __init__, forward

__init__(scale_factor)

Overview

Initialize the BilinearUpsample class.

Arguments: - scale_factor (:obj:Union[float, List[float]]): The multiplier for the spatial size.

forward(x)

Overview

Return the upsampled input tensor.

Arguments: - x (:obj:torch.Tensor): The input tensor. Returns: - upsample(:obj:torch.Tensor): The upsampled input tensor.

NoiseLinearLayer

Bases: Module

Overview

This is a linear layer with random noise.

Interfaces: __init__, reset_noise, reset_parameters, forward

__init__(in_channels, out_channels, sigma0=0.4)

Overview

Initialize the NoiseLinearLayer class. The 'enable_noise' attribute enables external control over whether noise is applied. - If enable_noise is True, the layer adds noise even if the module is in evaluation mode. - If enable_noise is False, no noise is added regardless of self.training.

Arguments: - in_channels (:obj:int): Number of channels in the input tensor. - out_channels (:obj:int): Number of channels in the output tensor. - sigma0 (:obj:int, optional): Default noise volume when initializing NoiseLinearLayer. Default is 0.4.

reset_noise()

Overview

Reset the noise settings in the layer.

reset_parameters()

Overview

Reset the parameters in the layer.

forward(x)

Overview

Perform the forward pass with noise.

Arguments: - x (:obj:torch.Tensor): The input tensor. Returns: - output (:obj:torch.Tensor): The output tensor with noise.

NaiveFlatten

Bases: Module

Overview

This module is a naive implementation of the flatten operation.

Interfaces: __init__, forward

__init__(start_dim=1, end_dim=-1)

Overview

Initialize the NaiveFlatten class.

Arguments: - start_dim (:obj:int, optional): The first dimension to flatten. Default is 1. - end_dim (:obj:int, optional): The last dimension to flatten. Default is -1.

forward(x)

Overview

Perform the flatten operation on the input tensor.

Arguments: - x (:obj:torch.Tensor): The input tensor. Returns: - output (:obj:torch.Tensor): The flattened output tensor.

weight_init_(weight, init_type='xavier', activation=None)

Overview

Initialize weight according to the specified type.

Arguments: - weight (:obj:torch.Tensor): The weight that needs to be initialized. - init_type (:obj:str, optional): The type of initialization to implement, supports ["xavier", "kaiming", "orthogonal"]. - activation (:obj:str, optional): The activation function name. Recommended to use only with ['relu', 'leaky_relu'].

sequential_pack(layers)

Overview

Pack the layers in the input list to a nn.Sequential module. If there is a convolutional layer in module, an extra attribute out_channels will be added to the module and set to the out_channel of the conv layer.

Arguments: - layers (:obj:List[nn.Module]): The input list of layers. Returns: - seq (:obj:nn.Sequential): Packed sequential container.

conv1d_block(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, activation=None, norm_type=None)

Overview

Create a 1-dimensional convolution layer with activation and normalization.

Arguments: - in_channels (:obj:int): Number of channels in the input tensor. - out_channels (:obj:int): Number of channels in the output tensor. - kernel_size (:obj:int): Size of the convolving kernel. - stride (:obj:int, optional): Stride of the convolution. Default is 1. - padding (:obj:int, optional): Zero-padding added to both sides of the input. Default is 0. - dilation (:obj:int, optional): Spacing between kernel elements. Default is 1. - groups (:obj:int, optional): Number of blocked connections from input channels to output channels. Default is 1. - activation (:obj:nn.Module, optional): The optional activation function. - norm_type (:obj:str, optional): Type of the normalization. Returns: - block (:obj:nn.Sequential): A sequential list containing the torch layers of the 1-dimensional convolution layer.

.. note:: Conv1d (https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html#torch.nn.Conv1d)

conv2d_block(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, pad_type='zero', activation=None, norm_type=None, num_groups_for_gn=1, bias=True)

Overview

Create a 2-dimensional convolution layer with activation and normalization.

Arguments: - in_channels (:obj:int): Number of channels in the input tensor. - out_channels (:obj:int): Number of channels in the output tensor. - kernel_size (:obj:int): Size of the convolving kernel. - stride (:obj:int, optional): Stride of the convolution. Default is 1. - padding (:obj:int, optional): Zero-padding added to both sides of the input. Default is 0. - dilation (:obj:int): Spacing between kernel elements. - groups (:obj:int, optional): Number of blocked connections from input channels to output channels. Default is 1. - pad_type (:obj:str, optional): The way to add padding, include ['zero', 'reflect', 'replicate']. Default is 'zero'. - activation (:obj:nn.Module): the optional activation function. - norm_type (:obj:str): The type of the normalization, now support ['BN', 'LN', 'IN', 'GN', 'SyncBN'], default set to None, which means no normalization. - num_groups_for_gn (:obj:int): Number of groups for GroupNorm. - bias (:obj:bool): whether to add a learnable bias to the nn.Conv2d. Default is True. Returns: - block (:obj:nn.Sequential): A sequential list containing the torch layers of the 2-dimensional convolution layer.

.. note:: Conv2d (https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d)

deconv2d_block(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, activation=None, norm_type=None)

Overview

Create a 2-dimensional transpose convolution layer with activation and normalization.

Arguments: - in_channels (:obj:int): Number of channels in the input tensor. - out_channels (:obj:int): Number of channels in the output tensor. - kernel_size (:obj:int): Size of the convolving kernel. - stride (:obj:int, optional): Stride of the convolution. Default is 1. - padding (:obj:int, optional): Zero-padding added to both sides of the input. Default is 0. - output_padding (:obj:int, optional): Additional size added to one side of the output shape. Default is 0. - groups (:obj:int, optional): Number of blocked connections from input channels to output channels. Default is 1. - activation (:obj:int, optional): The optional activation function. - norm_type (:obj:int, optional): Type of the normalization. Returns: - block (:obj:nn.Sequential): A sequential list containing the torch layers of the 2-dimensional transpose convolution layer.

.. note::

ConvTranspose2d (https://pytorch.org/docs/master/generated/torch.nn.ConvTranspose2d.html)

fc_block(in_channels, out_channels, activation=None, norm_type=None, use_dropout=False, dropout_probability=0.5)

Overview

Create a fully-connected block with activation, normalization, and dropout. Optional normalization can be done to the dim 1 (across the channels). x -> fc -> norm -> act -> dropout -> out

Arguments: - in_channels (:obj:int): Number of channels in the input tensor. - out_channels (:obj:int): Number of channels in the output tensor. - activation (:obj:nn.Module, optional): The optional activation function. - norm_type (:obj:str, optional): Type of the normalization. - use_dropout (:obj:bool, optional): Whether to use dropout in the fully-connected block. Default is False. - dropout_probability (:obj:float, optional): Probability of an element to be zeroed in the dropout. Default is 0.5. Returns: - block (:obj:nn.Sequential): A sequential list containing the torch layers of the fully-connected block.

.. note::

You can refer to nn.linear (https://pytorch.org/docs/master/generated/torch.nn.Linear.html).

normed_linear(in_features, out_features, bias=True, device=None, dtype=None, scale=1.0)

Overview

Create a nn.Linear module but with normalized fan-in init.

Arguments: - in_features (:obj:int): Number of features in the input tensor. - out_features (:obj:int): Number of features in the output tensor. - bias (:obj:bool, optional): Whether to add a learnable bias to the nn.Linear. Default is True. - device (:obj:torch.device, optional): The device to put the created module on. Default is None. - dtype (:obj:torch.dtype, optional): The desired data type of created module. Default is None. - scale (:obj:float, optional): The scale factor for initialization. Default is 1.0. Returns: - out (:obj:nn.Linear): A nn.Linear module with normalized fan-in init.

normed_conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', device=None, dtype=None, scale=1)

Overview

Create a nn.Conv2d module but with normalized fan-in init.

Arguments: - in_channels (:obj:int): Number of channels in the input tensor. - out_channels (:obj:int): Number of channels in the output tensor. - kernel_size (:obj:Union[int, Tuple[int, int]]): Size of the convolving kernel. - stride (:obj:Union[int, Tuple[int, int]], optional): Stride of the convolution. Default is 1. - padding (:obj:Union[int, Tuple[int, int]], optional): Zero-padding added to both sides of the input. Default is 0. - dilation (:Union[int, Tuple[int, int]], optional): Spacing between kernel elements. Default is 1. - groups (:obj:int, optional): Number of blocked connections from input channels to output channels. Default is 1. - bias (:obj:bool, optional): Whether to add a learnable bias to the nn.Conv2d. Default is True. - padding_mode (:obj:str, optional): The type of padding algorithm to use. Default is 'zeros'. - device (:obj:torch.device, optional): The device to put the created module on. Default is None. - dtype (:obj:torch.dtype, optional): The desired data type of created module. Default is None. - scale (:obj:float, optional): The scale factor for initialization. Default is 1. Returns: - out (:obj:nn.Conv2d): A nn.Conv2d module with normalized fan-in init.

MLP(in_channels, hidden_channels, out_channels, layer_num, layer_fn=None, activation=None, norm_type=None, use_dropout=False, dropout_probability=0.5, output_activation=True, output_norm=True, last_linear_layer_init_zero=False)

Overview

Create a multi-layer perceptron using fully-connected blocks with activation, normalization, and dropout, optional normalization can be done to the dim 1 (across the channels). x -> fc -> norm -> act -> dropout -> out

Arguments: - in_channels (:obj:int): Number of channels in the input tensor. - hidden_channels (:obj:int): Number of channels in the hidden tensor. - out_channels (:obj:int): Number of channels in the output tensor. - layer_num (:obj:int): Number of layers. - layer_fn (:obj:Callable, optional): Layer function. - activation (:obj:nn.Module, optional): The optional activation function. - norm_type (:obj:str, optional): The type of the normalization. - use_dropout (:obj:bool, optional): Whether to use dropout in the fully-connected block. Default is False. - dropout_probability (:obj:float, optional): Probability of an element to be zeroed in the dropout. Default is 0.5. - output_activation (:obj:bool, optional): Whether to use activation in the output layer. If True, we use the same activation as front layers. Default is True. - output_norm (:obj:bool, optional): Whether to use normalization in the output layer. If True, we use the same normalization as front layers. Default is True. - last_linear_layer_init_zero (:obj:bool, optional): Whether to use zero initializations for the last linear layer (including w and b), which can provide stable zero outputs in the beginning, usually used in the policy network in RL settings. Returns: - block (:obj:nn.Sequential): A sequential list containing the torch layers of the multi-layer perceptron.

.. note:: you can refer to nn.linear (https://pytorch.org/docs/master/generated/torch.nn.Linear.html).

one_hot(val, num, num_first=False)

Overview

Convert a torch.LongTensor to one-hot encoding. This implementation can be slightly faster than torch.nn.functional.one_hot.

Arguments: - val (:obj:torch.LongTensor): Each element contains the state to be encoded, the range should be [0, num-1] - num (:obj:int): Number of states of the one-hot encoding - num_first (:obj:bool, optional): If False, the one-hot encoding is added as the last dimension; otherwise, it is added as the first dimension. Default is False. Returns: - one_hot (:obj:torch.FloatTensor): The one-hot encoded tensor. Example: >>> one_hot(2torch.ones([2,2]).long(),3) tensor([[[0., 0., 1.], [0., 0., 1.]], [[0., 0., 1.], [0., 0., 1.]]]) >>> one_hot(2torch.ones([2,2]).long(),3,num_first=True) tensor([[[0., 0.], [1., 0.]], [[0., 1.], [0., 0.]], [[1., 0.], [0., 1.]]])

binary_encode(y, max_val)

Overview

Convert elements in a tensor to its binary representation.

Arguments: - y (:obj:torch.Tensor): The tensor to be converted into its binary representation. - max_val (:obj:torch.Tensor): The maximum value of the elements in the tensor. Returns: - binary (:obj:torch.Tensor): The input tensor in its binary representation. Example: >>> binary_encode(torch.tensor([3,2]),torch.tensor(8)) tensor([[0, 0, 1, 1],[0, 0, 1, 0]])

noise_block(in_channels, out_channels, activation=None, norm_type=None, use_dropout=False, dropout_probability=0.5, sigma0=0.4)

Overview

Create a fully-connected noise layer with activation, normalization, and dropout. Optional normalization can be done to the dim 1 (across the channels).

Arguments: - in_channels (:obj:int): Number of channels in the input tensor. - out_channels (:obj:int): Number of channels in the output tensor. - activation (:obj:str, optional): The optional activation function. Default is None. - norm_type (:obj:str, optional): Type of normalization. Default is None. - use_dropout (:obj:bool, optional): Whether to use dropout in the fully-connected block. - dropout_probability (:obj:float, optional): Probability of an element to be zeroed in the dropout. Default is 0.5. - sigma0 (:obj:float, optional): The sigma0 is the default noise volume when initializing NoiseLinearLayer. Default is 0.4. Returns: - block (:obj:nn.Sequential): A sequential list containing the torch layers of the fully-connected block.

Full Source Code

../ding/torch_utils/network/nn_module.py

1import math 2import torch 3import torch.nn as nn 4import torch.nn.functional as F 5from torch.nn.init import xavier_normal_, kaiming_normal_, orthogonal_ 6from typing import Union, Tuple, List, Callable 7from ding.compatibility import torch_ge_131 8 9from .normalization import build_normalization 10 11 12def weight_init_(weight: torch.Tensor, init_type: str = "xavier", activation: str = None) -> None: 13 """ 14 Overview: 15 Initialize weight according to the specified type. 16 Arguments: 17 - weight (:obj:`torch.Tensor`): The weight that needs to be initialized. 18 - init_type (:obj:`str`, optional): The type of initialization to implement, \ 19 supports ["xavier", "kaiming", "orthogonal"]. 20 - activation (:obj:`str`, optional): The activation function name. Recommended to use only with \ 21 ['relu', 'leaky_relu']. 22 """ 23 24 def xavier_init(weight, *args): 25 xavier_normal_(weight) 26 27 def kaiming_init(weight, activation): 28 assert activation is not None 29 if hasattr(activation, "negative_slope"): 30 kaiming_normal_(weight, a=activation.negative_slope) 31 else: 32 kaiming_normal_(weight, a=0) 33 34 def orthogonal_init(weight, *args): 35 orthogonal_(weight) 36 37 if init_type is None: 38 return 39 init_type_dict = {"xavier": xavier_init, "kaiming": kaiming_init, "orthogonal": orthogonal_init} 40 if init_type in init_type_dict: 41 init_type_dict[init_type](weight, activation) 42 else: 43 raise KeyError("Invalid Value in init type: {}".format(init_type)) 44 45 46def sequential_pack(layers: List[nn.Module]) -> nn.Sequential: 47 """ 48 Overview: 49 Pack the layers in the input list to a `nn.Sequential` module. 50 If there is a convolutional layer in module, an extra attribute `out_channels` will be added 51 to the module and set to the out_channel of the conv layer. 52 Arguments: 53 - layers (:obj:`List[nn.Module]`): The input list of layers. 54 Returns: 55 - seq (:obj:`nn.Sequential`): Packed sequential container. 56 """ 57 assert isinstance(layers, list) 58 seq = nn.Sequential(*layers) 59 for item in reversed(layers): 60 if isinstance(item, nn.Conv2d) or isinstance(item, nn.ConvTranspose2d): 61 seq.out_channels = item.out_channels 62 break 63 elif isinstance(item, nn.Conv1d): 64 seq.out_channels = item.out_channels 65 break 66 return seq 67 68 69def conv1d_block( 70 in_channels: int, 71 out_channels: int, 72 kernel_size: int, 73 stride: int = 1, 74 padding: int = 0, 75 dilation: int = 1, 76 groups: int = 1, 77 activation: nn.Module = None, 78 norm_type: str = None 79) -> nn.Sequential: 80 """ 81 Overview: 82 Create a 1-dimensional convolution layer with activation and normalization. 83 Arguments: 84 - in_channels (:obj:`int`): Number of channels in the input tensor. 85 - out_channels (:obj:`int`): Number of channels in the output tensor. 86 - kernel_size (:obj:`int`): Size of the convolving kernel. 87 - stride (:obj:`int`, optional): Stride of the convolution. Default is 1. 88 - padding (:obj:`int`, optional): Zero-padding added to both sides of the input. Default is 0. 89 - dilation (:obj:`int`, optional): Spacing between kernel elements. Default is 1. 90 - groups (:obj:`int`, optional): Number of blocked connections from input channels to output channels. \ 91 Default is 1. 92 - activation (:obj:`nn.Module`, optional): The optional activation function. 93 - norm_type (:obj:`str`, optional): Type of the normalization. 94 Returns: 95 - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the 1-dimensional \ 96 convolution layer. 97 98 .. note:: 99 Conv1d (https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html#torch.nn.Conv1d) 100 """ 101 block = [] 102 block.append(nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups)) 103 if norm_type is not None: 104 block.append(build_normalization(norm_type, dim=1)(out_channels)) 105 if activation is not None: 106 block.append(activation) 107 return sequential_pack(block) 108 109 110def conv2d_block( 111 in_channels: int, 112 out_channels: int, 113 kernel_size: int, 114 stride: int = 1, 115 padding: int = 0, 116 dilation: int = 1, 117 groups: int = 1, 118 pad_type: str = 'zero', 119 activation: nn.Module = None, 120 norm_type: str = None, 121 num_groups_for_gn: int = 1, 122 bias: bool = True 123) -> nn.Sequential: 124 """ 125 Overview: 126 Create a 2-dimensional convolution layer with activation and normalization. 127 Arguments: 128 - in_channels (:obj:`int`): Number of channels in the input tensor. 129 - out_channels (:obj:`int`): Number of channels in the output tensor. 130 - kernel_size (:obj:`int`): Size of the convolving kernel. 131 - stride (:obj:`int`, optional): Stride of the convolution. Default is 1. 132 - padding (:obj:`int`, optional): Zero-padding added to both sides of the input. Default is 0. 133 - dilation (:obj:`int`): Spacing between kernel elements. 134 - groups (:obj:`int`, optional): Number of blocked connections from input channels to output channels. \ 135 Default is 1. 136 - pad_type (:obj:`str`, optional): The way to add padding, include ['zero', 'reflect', 'replicate']. \ 137 Default is 'zero'. 138 - activation (:obj:`nn.Module`): the optional activation function. 139 - norm_type (:obj:`str`): The type of the normalization, now support ['BN', 'LN', 'IN', 'GN', 'SyncBN'], \ 140 default set to None, which means no normalization. 141 - num_groups_for_gn (:obj:`int`): Number of groups for GroupNorm. 142 - bias (:obj:`bool`): whether to add a learnable bias to the nn.Conv2d. Default is True. 143 Returns: 144 - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the 2-dimensional \ 145 convolution layer. 146 147 .. note:: 148 Conv2d (https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d) 149 """ 150 block = [] 151 assert pad_type in ['zero', 'reflect', 'replication'], "invalid padding type: {}".format(pad_type) 152 if pad_type == 'zero': 153 pass 154 elif pad_type == 'reflect': 155 block.append(nn.ReflectionPad2d(padding)) 156 padding = 0 157 elif pad_type == 'replication': 158 block.append(nn.ReplicationPad2d(padding)) 159 padding = 0 160 block.append( 161 nn.Conv2d( 162 in_channels, 163 out_channels, 164 kernel_size, 165 stride, 166 padding=padding, 167 dilation=dilation, 168 groups=groups, 169 bias=bias 170 ) 171 ) 172 if norm_type is not None: 173 if norm_type == 'LN': 174 # LN is implemented as GroupNorm with 1 group. 175 block.append(nn.GroupNorm(1, out_channels)) 176 elif norm_type == 'GN': 177 block.append(nn.GroupNorm(num_groups_for_gn, out_channels)) 178 elif norm_type in ['BN', 'IN', 'SyncBN']: 179 block.append(build_normalization(norm_type, dim=2)(out_channels)) 180 else: 181 raise KeyError( 182 "Invalid value in norm_type: {}. The valid norm_type are " 183 "BN, LN, IN, GN and SyncBN.".format(norm_type) 184 ) 185 186 if activation is not None: 187 block.append(activation) 188 return sequential_pack(block) 189 190 191def deconv2d_block( 192 in_channels: int, 193 out_channels: int, 194 kernel_size: int, 195 stride: int = 1, 196 padding: int = 0, 197 output_padding: int = 0, 198 groups: int = 1, 199 activation: int = None, 200 norm_type: int = None 201) -> nn.Sequential: 202 """ 203 Overview: 204 Create a 2-dimensional transpose convolution layer with activation and normalization. 205 Arguments: 206 - in_channels (:obj:`int`): Number of channels in the input tensor. 207 - out_channels (:obj:`int`): Number of channels in the output tensor. 208 - kernel_size (:obj:`int`): Size of the convolving kernel. 209 - stride (:obj:`int`, optional): Stride of the convolution. Default is 1. 210 - padding (:obj:`int`, optional): Zero-padding added to both sides of the input. Default is 0. 211 - output_padding (:obj:`int`, optional): Additional size added to one side of the output shape. Default is 0. 212 - groups (:obj:`int`, optional): Number of blocked connections from input channels to output channels. \ 213 Default is 1. 214 - activation (:obj:`int`, optional): The optional activation function. 215 - norm_type (:obj:`int`, optional): Type of the normalization. 216 Returns: 217 - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the 2-dimensional \ 218 transpose convolution layer. 219 220 .. note:: 221 222 ConvTranspose2d (https://pytorch.org/docs/master/generated/torch.nn.ConvTranspose2d.html) 223 """ 224 block = [ 225 nn.ConvTranspose2d( 226 in_channels=in_channels, 227 out_channels=out_channels, 228 kernel_size=kernel_size, 229 stride=stride, 230 padding=padding, 231 output_padding=output_padding, 232 groups=groups 233 ) 234 ] 235 if norm_type is not None: 236 block.append(build_normalization(norm_type, dim=2)(out_channels)) 237 if activation is not None: 238 block.append(activation) 239 return sequential_pack(block) 240 241 242def fc_block( 243 in_channels: int, 244 out_channels: int, 245 activation: nn.Module = None, 246 norm_type: str = None, 247 use_dropout: bool = False, 248 dropout_probability: float = 0.5 249) -> nn.Sequential: 250 """ 251 Overview: 252 Create a fully-connected block with activation, normalization, and dropout. 253 Optional normalization can be done to the dim 1 (across the channels). 254 x -> fc -> norm -> act -> dropout -> out 255 Arguments: 256 - in_channels (:obj:`int`): Number of channels in the input tensor. 257 - out_channels (:obj:`int`): Number of channels in the output tensor. 258 - activation (:obj:`nn.Module`, optional): The optional activation function. 259 - norm_type (:obj:`str`, optional): Type of the normalization. 260 - use_dropout (:obj:`bool`, optional): Whether to use dropout in the fully-connected block. Default is False. 261 - dropout_probability (:obj:`float`, optional): Probability of an element to be zeroed in the dropout. \ 262 Default is 0.5. 263 Returns: 264 - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the fully-connected block. 265 266 .. note:: 267 268 You can refer to nn.linear (https://pytorch.org/docs/master/generated/torch.nn.Linear.html). 269 """ 270 block = [] 271 block.append(nn.Linear(in_channels, out_channels)) 272 if norm_type is not None: 273 block.append(build_normalization(norm_type, dim=1)(out_channels)) 274 if activation is not None: 275 block.append(activation) 276 if use_dropout: 277 block.append(nn.Dropout(dropout_probability)) 278 return sequential_pack(block) 279 280 281def normed_linear( 282 in_features: int, 283 out_features: int, 284 bias: bool = True, 285 device=None, 286 dtype=None, 287 scale: float = 1.0 288) -> nn.Linear: 289 """ 290 Overview: 291 Create a nn.Linear module but with normalized fan-in init. 292 Arguments: 293 - in_features (:obj:`int`): Number of features in the input tensor. 294 - out_features (:obj:`int`): Number of features in the output tensor. 295 - bias (:obj:`bool`, optional): Whether to add a learnable bias to the nn.Linear. Default is True. 296 - device (:obj:`torch.device`, optional): The device to put the created module on. Default is None. 297 - dtype (:obj:`torch.dtype`, optional): The desired data type of created module. Default is None. 298 - scale (:obj:`float`, optional): The scale factor for initialization. Default is 1.0. 299 Returns: 300 - out (:obj:`nn.Linear`): A nn.Linear module with normalized fan-in init. 301 """ 302 303 out = nn.Linear(in_features, out_features, bias) 304 305 out.weight.data *= scale / out.weight.norm(dim=1, p=2, keepdim=True) 306 if bias: 307 out.bias.data.zero_() 308 return out 309 310 311def normed_conv2d( 312 in_channels: int, 313 out_channels: int, 314 kernel_size: Union[int, Tuple[int, int]], 315 stride: Union[int, Tuple[int, int]] = 1, 316 padding: Union[int, Tuple[int, int]] = 0, 317 dilation: Union[int, Tuple[int, int]] = 1, 318 groups: int = 1, 319 bias: bool = True, 320 padding_mode: str = 'zeros', 321 device=None, 322 dtype=None, 323 scale: float = 1 324) -> nn.Conv2d: 325 """ 326 Overview: 327 Create a nn.Conv2d module but with normalized fan-in init. 328 Arguments: 329 - in_channels (:obj:`int`): Number of channels in the input tensor. 330 - out_channels (:obj:`int`): Number of channels in the output tensor. 331 - kernel_size (:obj:`Union[int, Tuple[int, int]]`): Size of the convolving kernel. 332 - stride (:obj:`Union[int, Tuple[int, int]]`, optional): Stride of the convolution. Default is 1. 333 - padding (:obj:`Union[int, Tuple[int, int]]`, optional): Zero-padding added to both sides of the input. \ 334 Default is 0. 335 - dilation (:`Union[int, Tuple[int, int]]`, optional): Spacing between kernel elements. Default is 1. 336 - groups (:obj:`int`, optional): Number of blocked connections from input channels to output channels. \ 337 Default is 1. 338 - bias (:obj:`bool`, optional): Whether to add a learnable bias to the nn.Conv2d. Default is True. 339 - padding_mode (:obj:`str`, optional): The type of padding algorithm to use. Default is 'zeros'. 340 - device (:obj:`torch.device`, optional): The device to put the created module on. Default is None. 341 - dtype (:obj:`torch.dtype`, optional): The desired data type of created module. Default is None. 342 - scale (:obj:`float`, optional): The scale factor for initialization. Default is 1. 343 Returns: 344 - out (:obj:`nn.Conv2d`): A nn.Conv2d module with normalized fan-in init. 345 """ 346 347 out = nn.Conv2d( 348 in_channels, 349 out_channels, 350 kernel_size, 351 stride, 352 padding, 353 dilation, 354 groups, 355 bias, 356 padding_mode, 357 ) 358 out.weight.data *= scale / out.weight.norm(dim=(1, 2, 3), p=2, keepdim=True) 359 if bias: 360 out.bias.data.zero_() 361 return out 362 363 364def MLP( 365 in_channels: int, 366 hidden_channels: int, 367 out_channels: int, 368 layer_num: int, 369 layer_fn: Callable = None, 370 activation: nn.Module = None, 371 norm_type: str = None, 372 use_dropout: bool = False, 373 dropout_probability: float = 0.5, 374 output_activation: bool = True, 375 output_norm: bool = True, 376 last_linear_layer_init_zero: bool = False 377): 378 """ 379 Overview: 380 Create a multi-layer perceptron using fully-connected blocks with activation, normalization, and dropout, 381 optional normalization can be done to the dim 1 (across the channels). 382 x -> fc -> norm -> act -> dropout -> out 383 Arguments: 384 - in_channels (:obj:`int`): Number of channels in the input tensor. 385 - hidden_channels (:obj:`int`): Number of channels in the hidden tensor. 386 - out_channels (:obj:`int`): Number of channels in the output tensor. 387 - layer_num (:obj:`int`): Number of layers. 388 - layer_fn (:obj:`Callable`, optional): Layer function. 389 - activation (:obj:`nn.Module`, optional): The optional activation function. 390 - norm_type (:obj:`str`, optional): The type of the normalization. 391 - use_dropout (:obj:`bool`, optional): Whether to use dropout in the fully-connected block. Default is False. 392 - dropout_probability (:obj:`float`, optional): Probability of an element to be zeroed in the dropout. \ 393 Default is 0.5. 394 - output_activation (:obj:`bool`, optional): Whether to use activation in the output layer. If True, \ 395 we use the same activation as front layers. Default is True. 396 - output_norm (:obj:`bool`, optional): Whether to use normalization in the output layer. If True, \ 397 we use the same normalization as front layers. Default is True. 398 - last_linear_layer_init_zero (:obj:`bool`, optional): Whether to use zero initializations for the last \ 399 linear layer (including w and b), which can provide stable zero outputs in the beginning, \ 400 usually used in the policy network in RL settings. 401 Returns: 402 - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the multi-layer perceptron. 403 404 .. note:: 405 you can refer to nn.linear (https://pytorch.org/docs/master/generated/torch.nn.Linear.html). 406 """ 407 assert layer_num >= 0, layer_num 408 if layer_num == 0: 409 return sequential_pack([nn.Identity()]) 410 411 channels = [in_channels] + [hidden_channels] * (layer_num - 1) + [out_channels] 412 if layer_fn is None: 413 layer_fn = nn.Linear 414 block = [] 415 for i, (in_channels, out_channels) in enumerate(zip(channels[:-2], channels[1:-1])): 416 block.append(layer_fn(in_channels, out_channels)) 417 if norm_type is not None: 418 block.append(build_normalization(norm_type, dim=1)(out_channels)) 419 if activation is not None: 420 block.append(activation) 421 if use_dropout: 422 block.append(nn.Dropout(dropout_probability)) 423 424 # The last layer 425 in_channels = channels[-2] 426 out_channels = channels[-1] 427 block.append(layer_fn(in_channels, out_channels)) 428 """ 429 In the final layer of a neural network, whether to use normalization and activation are typically determined 430 based on user specifications. These specifications depend on the problem at hand and the desired properties of 431 the model's output. 432 """ 433 if output_norm is True: 434 # The last layer uses the same norm as front layers. 435 if norm_type is not None: 436 block.append(build_normalization(norm_type, dim=1)(out_channels)) 437 if output_activation is True: 438 # The last layer uses the same activation as front layers. 439 if activation is not None: 440 block.append(activation) 441 if use_dropout: 442 block.append(nn.Dropout(dropout_probability)) 443 444 if last_linear_layer_init_zero: 445 # Locate the last linear layer and initialize its weights and biases to 0. 446 for _, layer in enumerate(reversed(block)): 447 if isinstance(layer, nn.Linear): 448 nn.init.zeros_(layer.weight) 449 nn.init.zeros_(layer.bias) 450 break 451 452 return sequential_pack(block) 453 454 455class ChannelShuffle(nn.Module): 456 """ 457 Overview: 458 Apply channel shuffle to the input tensor. For more details about the channel shuffle, 459 please refer to the 'ShuffleNet' paper: https://arxiv.org/abs/1707.01083 460 Interfaces: 461 ``__init__``, ``forward`` 462 """ 463 464 def __init__(self, group_num: int) -> None: 465 """ 466 Overview: 467 Initialize the ChannelShuffle class. 468 Arguments: 469 - group_num (:obj:`int`): The number of groups to exchange. 470 """ 471 super().__init__() 472 self.group_num = group_num 473 474 def forward(self, x: torch.Tensor) -> torch.Tensor: 475 """ 476 Overview: 477 Forward pass through the ChannelShuffle module. 478 Arguments: 479 - x (:obj:`torch.Tensor`): The input tensor. 480 Returns: 481 - x (:obj:`torch.Tensor`): The shuffled input tensor. 482 """ 483 b, c, h, w = x.shape 484 g = self.group_num 485 assert (c % g == 0) 486 x = x.view(b, g, c // g, h, w).permute(0, 2, 1, 3, 4).contiguous().view(b, c, h, w) 487 return x 488 489 490def one_hot(val: torch.LongTensor, num: int, num_first: bool = False) -> torch.FloatTensor: 491 """ 492 Overview: 493 Convert a torch.LongTensor to one-hot encoding. This implementation can be slightly faster than 494 ``torch.nn.functional.one_hot``. 495 Arguments: 496 - val (:obj:`torch.LongTensor`): Each element contains the state to be encoded, the range should be [0, num-1] 497 - num (:obj:`int`): Number of states of the one-hot encoding 498 - num_first (:obj:`bool`, optional): If False, the one-hot encoding is added as the last dimension; otherwise, \ 499 it is added as the first dimension. Default is False. 500 Returns: 501 - one_hot (:obj:`torch.FloatTensor`): The one-hot encoded tensor. 502 Example: 503 >>> one_hot(2*torch.ones([2,2]).long(),3) 504 tensor([[[0., 0., 1.], 505 [0., 0., 1.]], 506 [[0., 0., 1.], 507 [0., 0., 1.]]]) 508 >>> one_hot(2*torch.ones([2,2]).long(),3,num_first=True) 509 tensor([[[0., 0.], [1., 0.]], 510 [[0., 1.], [0., 0.]], 511 [[1., 0.], [0., 1.]]]) 512 """ 513 assert (isinstance(val, torch.Tensor)), type(val) 514 assert val.dtype == torch.long 515 assert (len(val.shape) >= 1) 516 old_shape = val.shape 517 val_reshape = val.reshape(-1, 1) 518 ret = torch.zeros(val_reshape.shape[0], num, device=val.device) 519 # To remember the location where the original value is -1 in val. 520 # If the value is -1, then it should be converted to all zeros encodings and 521 # the corresponding entry in index_neg_one is 1, which is used to transform 522 # the ret after the operation of ret.scatter_(1, val_reshape, 1) to their correct encodings bellowing 523 index_neg_one = torch.eq(val_reshape, -1).float() 524 if index_neg_one.sum() != 0: # if -1 exists in val 525 # convert the original value -1 to 0 526 val_reshape = torch.where( 527 val_reshape != -1, val_reshape, 528 torch.zeros(val_reshape.shape, device=val.device).long() 529 ) 530 try: 531 ret.scatter_(1, val_reshape, 1) 532 if index_neg_one.sum() != 0: # if -1 exists in val 533 ret = ret * (1 - index_neg_one) # change -1's encoding from [1,0,...,0] to [0,0,...,0] 534 except RuntimeError: 535 raise RuntimeError('value: {}\nnum: {}\t:val_shape: {}\n'.format(val_reshape, num, val_reshape.shape)) 536 if num_first: 537 return ret.permute(1, 0).reshape(num, *old_shape) 538 else: 539 return ret.reshape(*old_shape, num) 540 541 542class NearestUpsample(nn.Module): 543 """ 544 Overview: 545 This module upsamples the input to the given scale_factor using the nearest mode. 546 Interfaces: 547 ``__init__``, ``forward`` 548 """ 549 550 def __init__(self, scale_factor: Union[float, List[float]]) -> None: 551 """ 552 Overview: 553 Initialize the NearestUpsample class. 554 Arguments: 555 - scale_factor (:obj:`Union[float, List[float]]`): The multiplier for the spatial size. 556 """ 557 super(NearestUpsample, self).__init__() 558 self.scale_factor = scale_factor 559 560 def forward(self, x: torch.Tensor) -> torch.Tensor: 561 """ 562 Overview: 563 Return the upsampled input tensor. 564 Arguments: 565 - x (:obj:`torch.Tensor`): The input tensor. 566 Returns: 567 - upsample(:obj:`torch.Tensor`): The upsampled input tensor. 568 """ 569 return F.interpolate(x, scale_factor=self.scale_factor, mode='nearest') 570 571 572class BilinearUpsample(nn.Module): 573 """ 574 Overview: 575 This module upsamples the input to the given scale_factor using the bilinear mode. 576 Interfaces: 577 ``__init__``, ``forward`` 578 """ 579 580 def __init__(self, scale_factor: Union[float, List[float]]) -> None: 581 """ 582 Overview: 583 Initialize the BilinearUpsample class. 584 Arguments: 585 - scale_factor (:obj:`Union[float, List[float]]`): The multiplier for the spatial size. 586 """ 587 super(BilinearUpsample, self).__init__() 588 self.scale_factor = scale_factor 589 590 def forward(self, x: torch.Tensor) -> torch.Tensor: 591 """ 592 Overview: 593 Return the upsampled input tensor. 594 Arguments: 595 - x (:obj:`torch.Tensor`): The input tensor. 596 Returns: 597 - upsample(:obj:`torch.Tensor`): The upsampled input tensor. 598 """ 599 return F.interpolate(x, scale_factor=self.scale_factor, mode='bilinear', align_corners=False) 600 601 602def binary_encode(y: torch.Tensor, max_val: torch.Tensor) -> torch.Tensor: 603 """ 604 Overview: 605 Convert elements in a tensor to its binary representation. 606 Arguments: 607 - y (:obj:`torch.Tensor`): The tensor to be converted into its binary representation. 608 - max_val (:obj:`torch.Tensor`): The maximum value of the elements in the tensor. 609 Returns: 610 - binary (:obj:`torch.Tensor`): The input tensor in its binary representation. 611 Example: 612 >>> binary_encode(torch.tensor([3,2]),torch.tensor(8)) 613 tensor([[0, 0, 1, 1],[0, 0, 1, 0]]) 614 """ 615 assert (max_val > 0) 616 x = y.clamp(0, max_val) 617 L = int(math.log(max_val, 2)) + 1 618 binary = [] 619 one = torch.ones_like(x) 620 zero = torch.zeros_like(x) 621 for i in range(L): 622 num = 1 << (L - i - 1) # 2**(L-i-1) 623 bit = torch.where(x >= num, one, zero) 624 x -= bit * num 625 binary.append(bit) 626 return torch.stack(binary, dim=1) 627 628 629class NoiseLinearLayer(nn.Module): 630 """ 631 Overview: 632 This is a linear layer with random noise. 633 Interfaces: 634 ``__init__``, ``reset_noise``, ``reset_parameters``, ``forward`` 635 """ 636 637 def __init__(self, in_channels: int, out_channels: int, sigma0: int = 0.4) -> None: 638 """ 639 Overview: 640 Initialize the NoiseLinearLayer class. The 'enable_noise' attribute enables external control over whether \ 641 noise is applied. 642 - If enable_noise is True, the layer adds noise even if the module is in evaluation mode. 643 - If enable_noise is False, no noise is added regardless of self.training. 644 Arguments: 645 - in_channels (:obj:`int`): Number of channels in the input tensor. 646 - out_channels (:obj:`int`): Number of channels in the output tensor. 647 - sigma0 (:obj:`int`, optional): Default noise volume when initializing NoiseLinearLayer. \ 648 Default is 0.4. 649 """ 650 super(NoiseLinearLayer, self).__init__() 651 self.in_channels = in_channels 652 self.out_channels = out_channels 653 self.weight_mu = nn.Parameter(torch.Tensor(out_channels, in_channels)) 654 self.weight_sigma = nn.Parameter(torch.Tensor(out_channels, in_channels)) 655 self.bias_mu = nn.Parameter(torch.Tensor(out_channels)) 656 self.bias_sigma = nn.Parameter(torch.Tensor(out_channels)) 657 self.register_buffer("weight_eps", torch.empty(out_channels, in_channels)) 658 self.register_buffer("bias_eps", torch.empty(out_channels)) 659 self.sigma0 = sigma0 660 self.enable_noise = False 661 self.reset_parameters() 662 self.reset_noise() 663 664 def _scale_noise(self, size: Union[int, Tuple]): 665 """ 666 Overview: 667 Scale the noise. 668 Arguments: 669 - size (:obj:`Union[int, Tuple]`): The size of the noise. 670 """ 671 672 x = torch.randn(size) 673 x = x.sign().mul(x.abs().sqrt()) 674 return x 675 676 def reset_noise(self): 677 """ 678 Overview: 679 Reset the noise settings in the layer. 680 """ 681 is_cuda = self.weight_mu.is_cuda 682 in_noise = self._scale_noise(self.in_channels).to(torch.device("cuda" if is_cuda else "cpu")) 683 out_noise = self._scale_noise(self.out_channels).to(torch.device("cuda" if is_cuda else "cpu")) 684 self.weight_eps = out_noise.ger(in_noise) 685 self.bias_eps = out_noise 686 687 def reset_parameters(self): 688 """ 689 Overview: 690 Reset the parameters in the layer. 691 """ 692 stdv = 1. / math.sqrt(self.in_channels) 693 self.weight_mu.data.uniform_(-stdv, stdv) 694 self.bias_mu.data.uniform_(-stdv, stdv) 695 696 std_weight = self.sigma0 / math.sqrt(self.in_channels) 697 self.weight_sigma.data.fill_(std_weight) 698 std_bias = self.sigma0 / math.sqrt(self.out_channels) 699 self.bias_sigma.data.fill_(std_bias) 700 701 def forward(self, x: torch.Tensor): 702 """ 703 Overview: 704 Perform the forward pass with noise. 705 Arguments: 706 - x (:obj:`torch.Tensor`): The input tensor. 707 Returns: 708 - output (:obj:`torch.Tensor`): The output tensor with noise. 709 """ 710 # Determine whether to add noise: 711 if self.enable_noise: 712 return F.linear( 713 x, 714 self.weight_mu + self.weight_sigma * self.weight_eps, 715 self.bias_mu + self.bias_sigma * self.bias_eps, 716 ) 717 else: 718 return F.linear(x, self.weight_mu, self.bias_mu) 719 720 721def noise_block( 722 in_channels: int, 723 out_channels: int, 724 activation: str = None, 725 norm_type: str = None, 726 use_dropout: bool = False, 727 dropout_probability: float = 0.5, 728 sigma0: float = 0.4 729): 730 """ 731 Overview: 732 Create a fully-connected noise layer with activation, normalization, and dropout. 733 Optional normalization can be done to the dim 1 (across the channels). 734 Arguments: 735 - in_channels (:obj:`int`): Number of channels in the input tensor. 736 - out_channels (:obj:`int`): Number of channels in the output tensor. 737 - activation (:obj:`str`, optional): The optional activation function. Default is None. 738 - norm_type (:obj:`str`, optional): Type of normalization. Default is None. 739 - use_dropout (:obj:`bool`, optional): Whether to use dropout in the fully-connected block. 740 - dropout_probability (:obj:`float`, optional): Probability of an element to be zeroed in the dropout. \ 741 Default is 0.5. 742 - sigma0 (:obj:`float`, optional): The sigma0 is the default noise volume when initializing NoiseLinearLayer. \ 743 Default is 0.4. 744 Returns: 745 - block (:obj:`nn.Sequential`): A sequential list containing the torch layers of the fully-connected block. 746 """ 747 block = [NoiseLinearLayer(in_channels, out_channels, sigma0=sigma0)] 748 if norm_type is not None: 749 block.append(build_normalization(norm_type, dim=1)(out_channels)) 750 if activation is not None: 751 block.append(activation) 752 if use_dropout: 753 block.append(nn.Dropout(dropout_probability)) 754 return sequential_pack(block) 755 756 757class NaiveFlatten(nn.Module): 758 """ 759 Overview: 760 This module is a naive implementation of the flatten operation. 761 Interfaces: 762 ``__init__``, ``forward`` 763 """ 764 765 def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None: 766 """ 767 Overview: 768 Initialize the NaiveFlatten class. 769 Arguments: 770 - start_dim (:obj:`int`, optional): The first dimension to flatten. Default is 1. 771 - end_dim (:obj:`int`, optional): The last dimension to flatten. Default is -1. 772 """ 773 super(NaiveFlatten, self).__init__() 774 self.start_dim = start_dim 775 self.end_dim = end_dim 776 777 def forward(self, x: torch.Tensor) -> torch.Tensor: 778 """ 779 Overview: 780 Perform the flatten operation on the input tensor. 781 Arguments: 782 - x (:obj:`torch.Tensor`): The input tensor. 783 Returns: 784 - output (:obj:`torch.Tensor`): The flattened output tensor. 785 """ 786 if self.end_dim != -1: 787 return x.view(*x.shape[:self.start_dim], -1, *x.shape[self.end_dim + 1:]) 788 else: 789 return x.view(*x.shape[:self.start_dim], -1) 790 791 792if torch_ge_131(): 793 Flatten = nn.Flatten 794else: 795 Flatten = NaiveFlatten