Python module
conv
The conv module provides classes for performing convolution operations in
various dimensions (1D, 2D, and 3D) on tensor inputs. These convolution
operations are core building blocks for neural networks, especially in computer
vision and sequence processing tasks.
Here’s an example demonstrating how to use a 1D convolution:
import max.nn as nn
from max.graph import Graph, ops, Weight, DeviceRef
from max.dtype import DType
import numpy as np
with Graph(name="conv_example") as graph:
# Define dimensions
batch_size = 2
seq_length = 10
in_channels = 16
out_channels = 32
kernel_size = 3
# Create input tensor [batch_size, sequence_length, channels]
x_data = np.zeros((batch_size, seq_length, in_channels), dtype=np.float32)
x = ops.constant(x_data, dtype=DType.float32, device=DeviceRef.CPU())
# Create weights for convolution
filter_1d = Weight(
name="filter_weight",
dtype=DType.float32,
shape=[kernel_size, in_channels, out_channels]
device=DeviceRef.CPU()
)
bias_1d = Weight(
name="bias_weight",
dtype=DType.float32,
shape=[out_channels]
device=DeviceRef.CPU()
)
# Create and apply Conv1D layer
conv1d = nn.Conv1D(
filter=filter_1d,
bias=bias_1d,
stride=1,
padding=1
)
output_1d = conv1d(x)
print(f"Conv1D output shape: {output_1d.shape}")
# Output: Conv1D output shape: [Dim(2), Dim(10), Dim(32)]Conv1D
class max.nn.legacy.conv.Conv1D(kernel_size, in_channels, out_channels, dtype, stride=1, padding=0, dilation=1, num_groups=1, device=None, has_bias=False, permute=False, name=None)
A 1D convolution over an input signal composed of several input planes.
Example:
conv = nn.Conv1D(
kernel_size=3,
in_channels=64,
out_channels=128,
dtype=DType.float32,
stride=1,
padding=0,
has_bias=False,
name="conv1d_weight",
device=DeviceRef.GPU(),
)-
Parameters:
bias
The optional bias vector stored on CPU with shape (out_channels,).
Model init moves the bias to device if present.
device
The device where matrix operations are performed.
dilation
dilation: int
Controls the dilation rate.
filter
filter: Weight
The weight matrix stored on CPU with shape (kernel_size, in_channels / num_groups, out_channels).
Model init moves the weight to device.
num_groups
num_groups: int
Number of blocked connections from input channels to output channels.
padding
Controls the amount of padding applied to the input.
If int: symmetric padding applied to both sides (pad_left = pad_right = padding). If tuple[int, int]: asymmetric padding as (pad_left, pad_right).
permute
permute: bool = False
bool controls whether self.filter is permuted from PyTorch order to max order. PyTorch order is: (out_channels, in_channels / num_groups, kernel_size) Max API order: (kernel_size, in_channels / num_groups, out_channels).
stride
stride: int
Controls the stride for the cross-correlation.
Conv2d
class max.nn.legacy.conv.Conv2d(kernel_size, in_channels, out_channels, dtype, stride=1, padding=0, dilation=1, num_groups=1, device=None, has_bias=False, permute=False, name=None)
A 2D convolution over an input signal composed of several input planes.
Example:
conv = nn.Conv2d(
kernel_size=3,
in_channels=64,
out_channels=128,
dtype=DType.float32,
stride=1,
padding=0,
has_bias=False,
name="conv2d_weight",
device=DeviceRef.GPU(),
)-
Parameters:
bias
The optional bias vector stored on CPU with shape (out_channels,).
Model init moves the bias to device if present.
device
The device where matrix operations are performed.
dilation
Controls the dilation rate.
filter
filter: Weight
The weight matrix stored on CPU with shape (height, width, in_channels / num_groups, out_channels).
Model init moves the weight to device.
num_groups
num_groups: int
Number of blocked connections from input channels to output channels.
padding
Controls the amount of padding applied before and after the input for height and width dimensions.
Format: (pad_top, pad_bottom, pad_left, pad_right).
permute
permute: bool = False
bool controls whether self.filter is permuted from PyTorch order to max order. PyTorch order is: (out_channels, in_channels / num_groups, height, width) Max API order: (height, width, in_channels / num_groups, out_channels).
shard()
shard(devices)
Creates sharded views of this Conv2d layer across multiple devices.
sharding_strategy
property sharding_strategy: ShardingStrategy | None
Get the Conv2d sharding strategy.
stride
Controls the stride for the cross-correlation.
Conv3D
class max.nn.legacy.conv.Conv3D(depth, height, width, in_channels, out_channels, dtype, stride=1, padding=0, dilation=1, num_groups=1, device=None, has_bias=False, permute=False, name=None)
A 3D convolution over an input signal composed of several input planes.
Example:
conv = nn.Conv3D(
depth=3,
height=3,
width=3,
in_channels=64,
out_channels=128,
dtype=DType.float32,
stride=1,
padding=0,
has_bias=False,
name="conv3d_weight",
device=DeviceRef.GPU(),
)-
Parameters:
bias
The optional bias vector stored on CPU with shape (out_channels,).
Model init moves the bias to device if present.
device
The device where matrix operations are performed.
dilation
Controls the dilation rate for depth, height, and width dimensions.
filter
filter: Weight
The weight matrix stored on CPU with shape (depth, height, width, in_channels / num_groups, out_channels).
Model init moves the weight to device.
num_groups
num_groups: int
Number of blocked connections from input channels to output channels.
padding
Controls the amount of padding applied before and after the input for depth, height, and width dimensions.
Format: (pad_front, pad_back, pad_top, pad_bottom, pad_left, pad_right).
permute
permute: bool = False
bool controls whether self.filter is permuted from PyTorch order to max order. PyTorch order is: (out_channels, in_channels / num_groups, depth, height, width) Max API order: (depth, height, width, in_channels / num_groups, out_channels).
stride
Controls the stride for the cross-correlation.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!