Upload folder using huggingface_hub

README.md

@@ -1,3 +1,73 @@
----
-license: apache-2.0
----
+# FaceNet Triplet ResNet Model (Grayscale, 112x112, Mobile-friendly)
+
+This repository provides a FaceNet-style triplet embedding model using ResNet backbones, optimized for mobile and edge devices:
+- Input: **Grayscale images** (`3` channel)
+- Resolution: **112x112 pixels**
+- Output: **Embeddings** suitable for face recognition and verification
+
+## Model Details
+
+- **Architecture:** ResNet50 with NdLinear
+- **Embedding Dimension:** 512
+- **Input:** 3x112x112 grayscale images (NCHW format)
+- **Exported weights:** `model.safetensors`
+- **Config:** `config.json`
+
+## Usage
+
+### 1. Clone or Download Files
+
+Download/copy the `models/` directory and dependencies (`ndlinear.py`, etc.) to your project.
+
+### 2. Install requirements
+
+```bash
+pip install torch safetensors
+```
+
+### 3. Load the model
+
+```python
+from models.resnet import Resnet50Triplet  # or your chosen variant
+
+model = Resnet50Triplet.from_pretrained(".", safe_serialization=True)
+model.eval()
+```
+
+### 4. Use for Face Recognition
+
+Obtain a face embedding from an input image, and compare embeddings (e.g., with cosine similarity) to recognize or verify identities.
+
+```python
+import torch
+
+# Example: batch of 1 grayscale image of 112x112
+images = torch.randn(1, 1, 112, 112)  # (batch_size, channels, height, width)
+
+with torch.no_grad():
+    embedding = model(images)  # embedding output suitable for face recognition
+print(embedding.shape)  # (batch_size, embedding_dim)
+```
+
+To perform recognition or verification, compare the embedding against a database of known face embeddings using distance/similarity metrics.
+
+## Files
+
+- `model.safetensors` - Model weights
+- `config.json` - Loader configuration
+- `models/` - Model definition files
+- `README.md` - This file
+
+## Notes
+
+- Model is optimized for runtime on edge/mobile devices (reduced input size, grayscale input for lower computational load)
+- Make sure your image preprocess pipeline produces three identical grayscaled channels, 112x112 images.
+
+## Credits
+
+- Backbone based on [PyTorch torchvision ResNet](https://pytorch.org/vision/stable/models/generated/torchvision.models.resnet50.html)
+- Architecture inspired by [Facenet PyTorch](https://github.com/timesler/facenet-pytorch)
+
+---
+
+*For contributions or issues, open a discussion or pull request.*

config.json

@@ -0,0 +1,5 @@
+{
+  "model_class": "Resnet50NdTriplet",
+  "embedding_dimension": 512,
+  "pretrained": false
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2bae12f646679e6b0b80acd9cd4cac9e2071d4c5e83053b04012b82787e14186
+size 94537792

models/ndlinear_util.py

@@ -0,0 +1,140 @@
+import json
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+import wandb
+from fvcore.nn import FlopCountAnalysis
+from sklearn.metrics import roc_curve
+from torchvision import models, transforms
+
+
+from ndlinear import NdLinear
+
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.RandomHorizontalFlip(),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2),
+    transforms.RandomRotation(10),
+    transforms.RandomResizedCrop((224, 224), scale=(0.8, 1.0)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
+])
+
+class ReshapedNdLinear(torch.nn.Module):
+    def __init__(self, nd_linear_layer):
+        super(ReshapedNdLinear, self).__init__()
+        self.nd_linear = nd_linear_layer
+
+    def forward(self, x):
+        x = x.reshape(*x.shape, 1)
+        x = self.nd_linear(x)
+        return x.view(x.size(0), -1)
+
+
+def print_cpu_layers(model):
+    found_cpu_layer = False
+    for name, module in model.named_modules():
+        if any(p.device.type == 'cpu' for p in module.parameters(recurse=False)):
+            print(f"Layer: {name}, Device: CPU")
+            found_cpu_layer = True
+    if not found_cpu_layer:
+        print("No layers are on the CPU.")
+
+
+def calculate_flops(model, input_tensor):
+    model.eval()
+    device = next(model.parameters()).device
+    input_tensor = input_tensor.to(device)
+    flops_analysis = FlopCountAnalysis(model, input_tensor)
+    flops = flops_analysis.total()
+    return flops
+
+
+def print_model_parameters(model):
+    return sum(p.numel() for p in model.parameters())
+
+
+def measure_latency_and_flops_cuda(model, input_tensor, warmup=10, runs=100):
+    assert torch.cuda.is_available(), "CUDA is not available."
+    device = torch.device('cuda')
+    model.to(device)
+    input_tensor = input_tensor.to(device)
+    model.eval()
+    torch.backends.cudnn.benchmark = True
+
+    with torch.no_grad():
+        for _ in range(warmup):
+            _ = model(input_tensor)
+        torch.cuda.synchronize()
+
+    timings = []
+    with torch.no_grad():
+        for _ in range(runs):
+            start = time.time()
+            _ = model(input_tensor)
+            torch.cuda.synchronize()
+            end = time.time()
+            timings.append(end - start)
+
+    avg_latency = sum(timings) / len(timings)
+    flops = calculate_flops(model, input_tensor[:1, ...])
+
+    print(f"Average CUDA Latency over {runs} runs: {avg_latency * 1000:.3f} ms")
+    print(f"Approx. FPS: {1.0 / avg_latency:.2f}")
+    print(f"Approx. Flops: {flops / 10 ** 9:.2f} GFlops")
+
+    return avg_latency, flops
+
+
+def modify_and_evaluate_backbone(model, cfg):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.train()
+
+    in_features = model.fc.in_features
+    fc_nd = NdLinear((in_features, 1), (cfg.embedding_size // 32, 32))
+    reshaped_fc = ReshapedNdLinear(fc_nd).to(device)
+
+    # Add dropout to the student model's fully connected layer
+    model.fc = nn.Sequential(
+        nn.Dropout(p=0.2),
+        reshaped_fc
+    )
+
+    for param in model.fc.parameters():
+        param.requires_grad = True
+
+    total_params = print_model_parameters(model)
+    wandb.log({"total_parameters": total_params})
+
+    model.to(device)
+    print_cpu_layers(model)
+    print(model)
+    return model
+
+
+def load_config(config_path='config.json'):
+    try:
+        with open(config_path, 'r') as f:
+            return json.load(f)
+    except FileNotFoundError as fe:
+        config = {
+            "learning_rate": 0.001,  # Adjusted learning rate
+            "epochs": 1000,
+            "batch_size": 32,
+            "eval_batch_size": 512,
+            "eval_every": 1000
+        }
+        return config
+
+
+def find_optimal_threshold(embeddings1, embeddings2, labels):
+    cosine_sim = np.sum(embeddings1 * embeddings2, axis=1)
+    fpr, tpr, thresholds = roc_curve(labels, cosine_sim)
+    # Youden's J statistic
+    j_scores = tpr - fpr
+    optimal_idx = np.argmax(j_scores)
+    optimal_threshold = thresholds[optimal_idx]
+    return optimal_threshold

models/resnet.py

@@ -0,0 +1,227 @@
+import torch.nn as nn
+from torch.nn import functional as F
+from ndlinear import NdLinear
+from .utils_resnet import resnet18, resnet34, resnet50, resnet101, resnet152, resnet34nd
+from .ndlinear_util import ReshapedNdLinear
+import torch
+from safetensors.torch import save_file as safe_save, load_file as safe_load
+import json
+import os
+
+class PretrainedTripletModel(nn.Module):
+    def save_pretrained(self, save_directory, safe_serialization=True):
+        os.makedirs(save_directory, exist_ok=True)
+        # Save config
+        config = {
+            "model_class": self.__class__.__name__,
+            "embedding_dimension": self.embedding_dimension,
+            "pretrained": self.pretrained,
+        }
+        with open(os.path.join(save_directory, "config.json"), "w") as f:
+            json.dump(config, f, indent=2)
+
+        # Save weights
+        state_dict = self.state_dict()
+        if safe_serialization:
+            safe_save(state_dict, os.path.join(save_directory, "model.safetensors"))
+        else:
+            torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin"))
+
+    @classmethod
+    def from_pretrained(cls, load_directory, safe_serialization=True):
+        # Load config
+        with open(os.path.join(load_directory, "config.json"), "r") as f:
+            config = json.load(f)
+        model = cls(**{k:config[k] for k in config if k in cls.__init__.__code__.co_varnames})
+        # Load weights
+        if safe_serialization:
+            state_dict = safe_load(os.path.join(load_directory, "model.safetensors"))
+        else:
+            state_dict = torch.load(os.path.join(load_directory, "pytorch_model.bin"), map_location="cpu")
+        model.load_state_dict(state_dict)
+        return model
+
+class Resnet18Triplet(nn.Module):
+    """Constructs a ResNet-18 model for FaceNet training using triplet loss.
+
+    Args:
+        embedding_dimension (int): Required dimension of the resulting embedding layer that is outputted by the model.
+                                   using triplet loss. Defaults to 512.
+        pretrained (bool): If True, returns a model pre-trained on the ImageNet dataset from a PyTorch repository.
+                           Defaults to False.
+    """
+
+    def __init__(self, embedding_dimension=512, pretrained=False):
+        super(Resnet18Triplet, self).__init__()
+        self.model = resnet18(pretrained=pretrained)
+
+        # Output embedding
+        input_features_fc_layer = self.model.fc.in_features
+        self.model.fc = nn.Linear(input_features_fc_layer, embedding_dimension, bias=False)
+
+    def forward(self, images):
+        """Forward pass to output the embedding vector (feature vector) after l2-normalization."""
+        embedding = self.model(images)
+        # From: https://github.com/timesler/facenet-pytorch/blob/master/models/inception_resnet_v1.py#L301
+        embedding = F.normalize(embedding, p=2, dim=1)
+
+        return embedding
+
+
+class Resnet34Triplet(nn.Module):
+    """Constructs a ResNet-34 model for FaceNet training using triplet loss.
+
+    Args:
+        embedding_dimension (int): Required dimension of the resulting embedding layer that is outputted by the model.
+                                   using triplet loss. Defaults to 512.
+        pretrained (bool): If True, returns a model pre-trained on the ImageNet dataset from a PyTorch repository.
+                           Defaults to False.
+    """
+
+    def __init__(self, embedding_dimension=512, pretrained=False):
+        super(Resnet34Triplet, self).__init__()
+        self.model = resnet34(pretrained=pretrained)
+
+        # Output embedding
+        input_features_fc_layer = self.model.fc.in_features
+        self.model.fc = nn.Linear(input_features_fc_layer, embedding_dimension, bias=False)
+
+    def forward(self, images):
+        """Forward pass to output the embedding vector (feature vector) after l2-normalization."""
+        embedding = self.model(images)
+        # From: https://github.com/timesler/facenet-pytorch/blob/master/models/inception_resnet_v1.py#L301
+        embedding = F.normalize(embedding, p=2, dim=1)
+
+        return embedding
+
+
+class Resnet50Triplet(nn.Module):
+    """Constructs a ResNet-50 model for FaceNet training using triplet loss.
+
+    Args:
+        embedding_dimension (int): Required dimension of the resulting embedding layer that is outputted by the model.
+                                   using triplet loss. Defaults to 512.
+        pretrained (bool): If True, returns a model pre-trained on the ImageNet dataset from a PyTorch repository.
+                           Defaults to False.
+    """
+
+    def __init__(self, embedding_dimension=512, pretrained=False):
+        super(Resnet50Triplet, self).__init__()
+        self.model = resnet50(pretrained=pretrained)
+
+        # Output embedding
+        input_features_fc_layer = self.model.fc.in_features
+        self.model.fc = nn.Linear(input_features_fc_layer, embedding_dimension, bias=False)
+
+    def forward(self, images):
+        """Forward pass to output the embedding vector (feature vector) after l2-normalization."""
+        embedding = self.model(images)
+        # From: https://github.com/timesler/facenet-pytorch/blob/master/models/inception_resnet_v1.py#L301
+        embedding = F.normalize(embedding, p=2, dim=1)
+
+        return embedding
+
+
+class Resnet101Triplet(nn.Module):
+    """Constructs a ResNet-101 model for FaceNet training using triplet loss.
+
+    Args:
+        embedding_dimension (int): Required dimension of the resulting embedding layer that is outputted by the model.
+                                   using triplet loss. Defaults to 512.
+        pretrained (bool): If True, returns a model pre-trained on the ImageNet dataset from a PyTorch repository.
+                           Defaults to False.
+    """
+
+    def __init__(self, embedding_dimension=512, pretrained=False):
+        super(Resnet101Triplet, self).__init__()
+        self.model = resnet101(pretrained=pretrained)
+
+        # Output embedding
+        input_features_fc_layer = self.model.fc.in_features
+        self.model.fc = nn.Linear(input_features_fc_layer, embedding_dimension, bias=False)
+
+    def forward(self, images):
+        """Forward pass to output the embedding vector (feature vector) after l2-normalization."""
+        embedding = self.model(images)
+        # From: https://github.com/timesler/facenet-pytorch/blob/master/models/inception_resnet_v1.py#L301
+        embedding = F.normalize(embedding, p=2, dim=1)
+
+        return embedding
+
+
+class Resnet152Triplet(nn.Module):
+    """Constructs a ResNet-152 model for FaceNet training using triplet loss.
+
+    Args:
+        embedding_dimension (int): Required dimension of the resulting embedding layer that is outputted by the model.
+                                   using triplet loss. Defaults to 512.
+        pretrained (bool): If True, returns a model pre-trained on the ImageNet dataset from a PyTorch repository.
+                           Defaults to False.
+    """
+
+    def __init__(self, embedding_dimension=512, pretrained=False):
+        super(Resnet152Triplet, self).__init__()
+        self.model = resnet152(pretrained=pretrained)
+
+        # Output embedding
+        input_features_fc_layer = self.model.fc.in_features
+        self.model.fc = nn.Linear(input_features_fc_layer, embedding_dimension, bias=False)
+
+    def forward(self, images):
+        """Forward pass to output the embedding vector (feature vector) after l2-normalization."""
+        embedding = self.model(images)
+        # From: https://github.com/timesler/facenet-pytorch/blob/master/models/inception_resnet_v1.py#L301
+        embedding = F.normalize(embedding, p=2, dim=1)
+
+        return embedding
+
+
+class Resnet34NdTriplet(nn.Module):
+    """Constructs a ResNet-34 model for FaceNet training using triplet loss.
+
+    Args:
+        embedding_dimension (int): Required dimension of the resulting embedding layer that is outputted by the model.
+                                   using triplet loss. Defaults to 512.
+        pretrained (bool): If True, returns a model pre-trained on the ImageNet dataset from a PyTorch repository.
+                           Defaults to False.
+    """
+
+    def __init__(self, embedding_dimension=512, pretrained=False):
+        super(Resnet34NdTriplet, self).__init__()
+        # self.model = resnet34nd(pretrained=False)
+        self.model = resnet34(pretrained=pretrained)
+
+        # Output embedding
+        input_features_fc_layer = self.model.fc.in_features
+        self.model.fc = ReshapedNdLinear(
+            NdLinear((input_features_fc_layer, 1),
+                     (embedding_dimension, 1),
+                     bias=False)
+        )
+
+    def forward(self, images):
+        """Forward pass to output the embedding vector (feature vector) after l2-normalization."""
+        embedding = self.model(images)
+        # From: https://github.com/timesler/facenet-pytorch/blob/master/models/inception_resnet_v1.py#L301
+        embedding = F.normalize(embedding, p=2, dim=1)
+
+        return embedding
+
+
+class Resnet50NdTriplet(PretrainedTripletModel):
+    def __init__(self, embedding_dimension=512, pretrained=False):
+        super().__init__()
+        self.embedding_dimension = embedding_dimension
+        self.pretrained = pretrained
+        self.model = resnet50(pretrained=pretrained)
+        input_features_fc_layer = self.model.fc.in_features
+        self.model.fc = ReshapedNdLinear(
+            NdLinear((input_features_fc_layer, 1), (embedding_dimension // 16, 16), bias=False)
+        )
+    def forward(self, images):
+        """Forward pass to output the embedding vector (feature vector) after l2-normalization."""
+        embedding = self.model(images)
+        # From: https://github.com/timesler/facenet-pytorch/blob/master/models/inception_resnet_v1.py#L301
+        embedding = F.normalize(embedding, p=2, dim=1)
+
+        return embedding

models/utils_resnet.py

@@ -0,0 +1,382 @@
+"""This code was imported from the official PyTorch Torchvision GitHub repository for the purposes doing experiments
+    with fine-tuned resnet architectures:
+    https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
+"""
+
+import torch
+import torch.nn as nn
+from ndlinear import NdLinear
+from .ndlinear_util import ReshapedNdLinear
+
+# Replaced 'from .utils import load_state_dict_from_url' with the following:
+try:
+    from torch.hub import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet34nd', 'resnet50', 'resnet50nd',
+           'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
+           'wide_resnet50_2', 'wide_resnet101_2']
+
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None):
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+                            self.base_width, previous_dilation, norm_layer))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x):
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+
+def _resnet(arch, block, layers, pretrained, progress, **kwargs):
+    model = ResNet(block, layers, **kwargs)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls[arch],
+                                              progress=progress)
+        model.load_state_dict(state_dict)
+    return model
+
+
+def resnet18(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-18 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
+                   **kwargs)
+
+
+def resnet34(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-34 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet50(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-50 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet101(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-101 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnet152(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-152 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
+                   **kwargs)
+
+
+def resnext50_32x4d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-50 32x4d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 4
+    return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def resnext101_32x8d(pretrained=False, progress=True, **kwargs):
+    r"""ResNeXt-101 32x8d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 8
+    return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-50-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
+                   pretrained, progress, **kwargs)
+
+
+def wide_resnet101_2(pretrained=False, progress=True, **kwargs):
+    r"""Wide ResNet-101-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
+                   pretrained, progress, **kwargs)
+
+
+def resnet34nd(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-34 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    _resnet34 = _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs)
+    input_dim = _resnet34.fc.in_features
+    output_dim = _resnet34.fc.out_features
+    _resnet34.fc = ReshapedNdLinear(NdLinear((input_dim, 1), (output_dim // 16, 16)))
+    return _resnet34
+
+
+def resnet50nd(pretrained=False, progress=True, **kwargs):
+    r"""ResNet-50 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    _resnet50 = _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
+    input_dim = _resnet50.fc.in_features
+    output_dim = _resnet50.fc.out_features
+    _resnet50.fc = ReshapedNdLinear(NdLinear((input_dim, 1), (output_dim // 16, 16)))
+    return _resnet50

ndlinear.py

@@ -0,0 +1,82 @@
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+
+class NdLinear(nn.Module):
+    def __init__(self, input_dims: tuple, hidden_size: tuple, transform_outer=True, bias=True):
+        """
+        NdLinear: A PyTorch layer for projecting tensors into multi-space representations.
+
+        Unlike conventional embedding layers that map into a single vector space, NdLinear
+        transforms tensors across a collection of vector spaces, capturing multivariate structure
+        and topical information that standard deep learning architectures typically lose.
+
+        Args:
+            input_dims (tuple): Shape of input tensor (excluding batch dimension).
+            hidden_size (tuple): Target hidden dimensions after transformation.
+        """
+        super(NdLinear, self).__init__()
+
+        if len(input_dims) != len(hidden_size):
+            raise Exception("Input shape and hidden shape do not match.")
+
+        self.input_dims = input_dims
+        self.hidden_size = hidden_size
+        self.num_layers = len(input_dims)  # Must match since dims are equal
+        # self.relu = nn.ReLU()
+        self.transform_outer = transform_outer
+
+        # Define transformation layers per dimension
+        self.align_layers = nn.ModuleList([
+            nn.Linear(input_dims[i], hidden_size[i], bias=bias) for i in range(self.num_layers)
+        ])
+
+    def forward(self, X):
+        """
+        Forward pass to project input tensor into a new multi-space representation.
+        - Incrementally transposes, flattens, applies linear layers, and restores shape.
+
+        Expected Input Shape: [batch_size, *input_dims]
+        Output Shape: [batch_size, *hidden_size]
+
+        Args:
+            X (torch.Tensor): Input tensor with shape [batch_size, *input_dims]
+
+        Returns:
+            torch.Tensor: Output tensor with shape [batch_size, *hidden_size]
+        """
+        num_transforms = self.num_layers  # Number of transformations
+
+        # Define iteration order
+        # transform_indices = range(num_transforms) if transform_outer else reversed(range(num_transforms))
+
+        for i in range(num_transforms):
+            if self.transform_outer:
+                layer = self.align_layers[i]
+                transpose_dim = i + 1
+            else:
+                layer = self.align_layers[num_transforms - (i+1)]
+                transpose_dim = num_transforms - i
+
+            # Transpose the selected dimension to the last position
+            X = torch.transpose(X, transpose_dim, num_transforms).contiguous()
+
+            # Store original shape before transformation
+            X_size = X.shape[:-1]
+
+            # Flatten everything except the last dimension
+            X = X.view(-1, X.shape[-1])
+
+            # Apply transformation
+            X = layer(X)
+
+            # Reshape back to the original spatial structure (with new embedding dim)
+            X = X.view(*X_size, X.shape[-1])
+
+            # Transpose the dimension back to its original position
+            X = torch.transpose(X, transpose_dim, num_transforms).contiguous()
+
+        return X
+