fovi.arch.knnalexnet

class fovi.arch.knnalexnet.KNNAlexNetBlock(in_channels, out_channels, k, fov, cmf_a, in_res, stride, style='isotropic', conv_layer='default', cart_res=None, pool=True, pool_mode='max', pool_k=9, pool_stride=2, norm_type='batch', arch_flag='', activation=nn.ReLU, sample_cortex=True, device='cuda', auto_match_cart_resources=0, ref_frame_mult=None)

Bases: Module

A block combining KNN convolution, optional pooling, normalization and nonlinearity.

This module implements a complete processing block that can include: - KNN convolution layer - Normalization (batch, layer, etc.) - Activation function - Optional pooling layer

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of output channels.

• k (int) – Number of nearest neighbors for KNN operations.

• fov (float) – Field of view parameter.

• cmf_a (float) – Cortical magnification factor parameter.

• in_res (int) – Input resolution.

• stride (int) – Stride for the convolution operation.

• style (str, optional) – Style of coordinate system. Defaults to ‘isotropic’.

• conv_layer (class or str, optional) – Convolution layer class or name. Defaults to ‘default’.

• ref_frame_mult (float, optional) – Multiplier for reference frame size. If not None, ref_frame_side_length = ceil(ref_frame_mult * sqrt(k)). Defaults to None.

• cart_res (int, optional) – Cartesian resolution. Defaults to None.

• pool (bool, optional) – Whether to include pooling layer. Defaults to True.

• pool_mode (str, optional) – Pooling mode (‘max’, ‘avg’). Defaults to ‘max’.

• pool_k (int, optional) – Number of neighbors for pooling. Defaults to 9.

• pool_stride (int, optional) – Stride for pooling. Defaults to 2.

• norm_type (str, optional) – Type of normalization. Defaults to ‘batch’.

• arch_flag (str, optional) – Architecture flag. Defaults to ‘’.

• activation (class, optional) – Activation function class. Defaults to nn.ReLU.

• sample_cortex (bool, optional) – Whether to sample cortex. Defaults to True.

• device (str, optional) – Device to use. Defaults to ‘cuda’.

• auto_match_cart_resources (int, optional) – Auto-match cartesian resources. Defaults to 0.

__init__(in_channels, out_channels, k, fov, cmf_a, in_res, stride, style='isotropic', conv_layer='default', cart_res=None, pool=True, pool_mode='max', pool_k=9, pool_stride=2, norm_type='batch', arch_flag='', activation=nn.ReLU, sample_cortex=True, device='cuda', auto_match_cart_resources=0, ref_frame_mult=None)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Forward pass through the block.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor after processing through conv, norm, activation, and pooling.

Return type:

torch.Tensor

class fovi.arch.knnalexnet.KNNAlexNet(in_res: int, in_channels: int, features_per_layer: list[int], stride_per_layer: list[int], pool_after: list[int], k_per_layer: list[int], n_classes: int = None, out_res: int = None, auto_match_cart_resources=0, style='isotropic', layer_cls='default', norm_type='batch', arch_flag='', fov=16, cmf_a=0.5, device='cuda', sample_cortex=True, ref_frame_mult=None)

Bases: Module

KNN-based AlexNet architecture with foveated processing.

This class implements a complete KNN-based AlexNet architecture that processes input through multiple KNN convolution blocks with optional pooling layers.

Parameters:

• in_res (int) – Input resolution.

• in_channels (int) – Number of input channels.

• features_per_layer (list[int]) – List of feature dimensions for each layer.

• stride_per_layer (list[int]) – List of stride values for each layer.

• pool_after (list[int]) – List of layer indices after which to apply pooling.

• k_per_layer (list[int]) – List of k values (nearest neighbors) for each layer.

• n_classes (int, optional) – Number of output classes. If None, no classifier is added. Defaults to None.

• out_res (int, optional) – Output resolution. If None, no output pooling is applied. Defaults to None.

• auto_match_cart_resources (int, optional) – Auto-match cartesian resources. Defaults to 0.

• style (str, optional) – Style of coordinate system. Defaults to ‘isotropic’.

• layer_cls (class or str, optional) – Convolution layer class or name. Defaults to ‘KNNConvLayer’.

• norm_type (str, optional) – Type of normalization. Defaults to ‘batch’.

• arch_flag (str, optional) – Architecture flag. Defaults to ‘’.

• fov (float, optional) – Field of view parameter. Defaults to 16.

• cmf_a (float, optional) – Cortical magnification factor parameter. Defaults to 0.5.

• device (str, optional) – Device to use. Defaults to ‘cuda’.

• sample_cortex (bool, optional) – Whether to sample cortex. Defaults to True.

• ref_frame_mult (float, optional) – Multiplier for reference frame size. If not None, ref_frame_side_length = ceil(ref_frame_mult * sqrt(k)). Defaults to None.

__init__(in_res: int, in_channels: int, features_per_layer: list[int], stride_per_layer: list[int], pool_after: list[int], k_per_layer: list[int], n_classes: int = None, out_res: int = None, auto_match_cart_resources=0, style='isotropic', layer_cls='default', norm_type='batch', arch_flag='', fov=16, cmf_a=0.5, device='cuda', sample_cortex=True, ref_frame_mult=None)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Forward pass through the entire network.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, in_channels, in_res, in_res).

Returns:

Output tensor. If n_classes is specified, returns logits of shape

(batch_size, n_classes). Otherwise, returns features of shape (batch_size, out_channels, out_coords.shape[0]).

Return type:

torch.Tensor