models

The models module contains neural network architectures used in Nenya for self-supervised learning.

Submodules

nenya.models.resnet_big

This module implements the ResNet backbone with projection head for contrastive learning.

Classes

class nenya.models.resnet_big.SupConResNet(name='resnet50', head='mlp', feat_dim=128)[source]

Base encoder network with projection head for contrastive learning.

Parameters:

  • name (str, optional) – Name of the backbone architecture (‘resnet18’, ‘resnet34’, ‘resnet50’, etc.)

  • head (str, optional) – Type of projection head (‘linear’, ‘mlp’)

  • feat_dim (int, optional) – Dimension of the feature vector

forward(x)[source]

Forward pass through the network.

Parameters:

x (torch.Tensor) – Input tensor

Returns:

Feature vector

Return type:

torch.Tensor

forward_feat(x)

Extract features without the projection head.

Parameters:

x (torch.Tensor) – Input tensor

Returns:

Feature vector before projection

Return type:

torch.Tensor

class nenya.models.resnet_big.SupCEResNet(name='resnet50', num_classes=10, head='mlp', feat_dim=128)[source]

Network for supervised contrastive learning with classification head.

Parameters:

  • name (str, optional) – Name of the backbone architecture

  • num_classes (int, optional) – Number of output classes

  • head (str, optional) – Type of projection head

  • feat_dim (int, optional) – Dimension of the feature vector

forward(x)[source]

Forward pass through the network.

Parameters:

x (torch.Tensor) – Input tensor

Returns:

Classification logits

Return type:

torch.Tensor

features(x)

Extract features before the classification head.

Parameters:

x (torch.Tensor) – Input tensor

Returns:

Feature vector

Return type:

torch.Tensor

Helper Functions

nenya.models.resnet_big.get_resnet(name, pretrained=False)

Get a ResNet model with a specific architecture.

Parameters:

  • name (str) – Name of the ResNet architecture (‘resnet18’, ‘resnet34’, ‘resnet50’, etc.)

  • pretrained (bool, optional) – Whether to use pre-trained weights. Defaults to False.

Returns:

ResNet model

Return type:

torch.nn.Module

Raises:

ValueError – If architecture name is not recognized

Architecture Details

ResNet Backbone

The models use ResNet architectures with varying depths:

  • ResNet18: 18 layers, ~11M parameters

  • ResNet34: 34 layers, ~21M parameters

  • ResNet50: 50 layers, ~23M parameters (default)

  • ResNet101: 101 layers, ~42M parameters

  • ResNet152: 152 layers, ~58M parameters

Projection Head

For contrastive learning, a projection head is added on top of the backbone:

  • Linear head: Single linear layer

  • MLP head: Two-layer MLP with ReLU activation (default)

The projection head maps the backbone features to a lower-dimensional space (typically 128 dimensions) where the contrastive loss is applied.

Usage Examples

Creating a model:

from nenya.models.resnet_big import SupConResNet

# Create a ResNet50 model with 128-dimensional features
model = SupConResNet(name='resnet50', feat_dim=128)

# Forward pass
import torch
x = torch.randn(10, 3, 64, 64)  # Batch of 10 images
features = model(x)  # Shape: [10, 128]

Model Input/Output

  • Input: Images with shape [batch_size, 3, height, width]

  • Output: Feature vectors with shape [batch_size, feat_dim]

For the contrastive learning setup:

  • Input is augmented pairs of images: [2*batch_size, 3, height, width]

  • Features are split and reshaped: [batch_size, 2, feat_dim]