bitorch_engine.layers.qconv.nbit.layer.nBitConv2dBase

class bitorch_engine.layers.qconv.nbit.layer.nBitConv2dBase(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, a_bit: int = 4, w_bit: int = 4, device=None, dtype=torch.float32)

Methods

__init__	Initializes the nBitConv2dBase module, a base class for creating convolutional layers with n-bit quantized weights.
generate_quantized_weight	Generates and sets the quantized weight based on the current weight parameter.
prepare_params	Prepares and initializes the model parameters for training.
reset_parameters	Initializes or resets the weight parameter using Kaiming uniform initialization.
set_quantized_weight_data	Sets the quantized weight parameter with the provided tensor.
set_weight_data	Sets the weight parameter with the provided tensor.

Attributes

opt_weight	Returns the proper weight parameter for the forward pass.
training

__init__(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, a_bit: int = 4, w_bit: int = 4, device=None, dtype=torch.float32)

Initializes the nBitConv2dBase module, a base class for creating convolutional layers with n-bit quantized weights.

Parameters:

• in_channels (int) – The number of input channels in the convolutional layer.

• out_channels (int) – The number of output channels in the convolutional layer.

• kernel_size (int) – The size of the convolutional kernel.

• stride (int, optional) – The stride of the convolution. Defaults to 1.

• padding (int, optional) – The padding added to all sides of the input tensor. Defaults to 0.

• dilation (int, optional) – The spacing between kernel elements. Defaults to 1.

• a_bit (int, optional) – The bit-width for activation quantization. Defaults to 4.

• w_bit (int, optional) – The bit-width for weight quantization. Defaults to 4.

• device (optional) – The device on which the module will be allocated. Defaults to None.

• dtype (optional) – The desired data type of the parameters. Defaults to torch.float.

generate_quantized_weight(qweight_only: bool = False) → None

Generates and sets the quantized weight based on the current weight parameter. This method should be overridden by subclasses to implement specific quantization logic.

Parameters:

qweight_only (bool, optional) – If True, the original weight tensor is discarded to save memory.

property opt_weight

Returns the proper weight parameter for the forward pass. If the model is in evaluation mode and quantized weights are available, it returns the quantized weights; otherwise, it returns the original weights.

Returns:

The optimal weight parameter for the forward pass.

Return type:

torch.nn.Parameter

prepare_params() → None

Prepares and initializes the model parameters for training.

Note

This method MUST be called after model initialization and before training starts to ensure the weights are properly prepared for efficient computation.

reset_parameters() → None

Initializes or resets the weight parameter using Kaiming uniform initialization.

set_quantized_weight_data(x: Tensor) → None

Sets the quantized weight parameter with the provided tensor.

Parameters:

x (torch.Tensor) – The tensor to be used as the new quantized weight parameter.

set_weight_data(x: Tensor) → None

Sets the weight parameter with the provided tensor.

Parameters:

x (torch.Tensor) – The tensor to be used as the new weight parameter.