ttnn.prepare_conv_weights

ttnn.prepare_conv_weights(*args, **kwargs)

TTNN Conv2D applies preprocessing to the weights tensors before performing the convolution operation, to convert the weights into a format suitable for the operation. This can be applied just once to the weights and bias tensors, and the resulting tensors can be reused for multiple invocations of the same convolution operation. The exact format of the weights and bias tensors depends on the input tensor parameters and the sharding scheme.

Parameters:

  • weight_tensor (ttnn.Tensor) – the weight tensor in PyTorch Conv2d format.

  • input_memory_config (ttnn.MemoryConfig) – the memory configuration for the input tensor.

  • input_layout (ttnn.Tensor) – the layout of the input tensor.

  • weights_format (ttnn.Tensor) – the format of the weights tensor. Currently only supports OIHW. (out_channels, in_channels, kernel_height, kernel_width)

  • in_channels (int) – number of input channels.

  • out_channels (int) – number of output channels.

  • batch_size (int) – batch size.

  • input_height (int) – height of the input tensor.

  • input_width (int) – width of the input tensor.

  • kernel_size (tuple[int, int]) – size of the convolving kernel.

  • stride (tuple[int, int]) – stride of the cross-correlation.

  • padding (tuple[int, int] or tuple[int, int, int, int])) – zero-padding added to both sides of the input. [pad_height, pad_width] or [pad_top, pad_bottom, pad_left, pad_right].

  • dilation (tuple[int, int]) – spacing between kernel elements.

  • has_bias (bool) – whether the convolution has a bias term.

  • groups (int) – number of blocked connections from input channels to output channels.

  • conv_config (ttnn.Conv2dConfig, None) – configuration for convolution. Default: None

  • compute_config (ttnn.DeviceComputeKernelConfig, None) – configuration for compute kernel. Default: None

Returns:

The preprocessed weight tensor on device

Return type:

[ttnn.Tensor]: The preprocessed bias tensor on device