depthwise_conv2d

depthwise_conv2d

Module for depth-wise 2D Convolution layer.

DepthwiseConv2D

DepthwiseConv2D(dimension, kernel_size, stride, padding)

Bases: Module

A depth-wise 2D convolution applies a single convolutional filter to each input channel independently, without mixing channels (unlike standard convolution).

Source code in src/versatil/models/layers/convolution/depthwise_conv2d.py

def __init__(self, dimension, kernel_size, stride, padding):
    super().__init__()
    self.convolution = nn.Conv2d(
        dimension, dimension, kernel_size, stride, padding, groups=dimension
    )

forward

forward(x)

Parameters:

Name	Type	Description	Default
x	Tensor	input of shape (batch, height, width, channels).	required

Returns:

Name	Type	Description
output	Tensor	tensor after applying depth-wise spatial convolutions, shape (batch, height, width, channels).

Source code in src/versatil/models/layers/convolution/depthwise_conv2d.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Args:
        x: input of shape (batch, height, width, channels).

    Returns:
        output: tensor after applying depth-wise spatial convolutions, shape (batch, height, width, channels).
    """
    x = x.permute(0, 3, 1, 2)
    x = self.convolution(x)
    x = x.permute(0, 2, 3, 1)
    return x