### Install submitit for Multi-node Training

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Installs the submitit library, which is used for managing multi-node training on SLURM clusters. This is a prerequisite for running the provided multi-node training scripts.

```bash
pip install submitit
```

--------------------------------

### Implement ConvNeXtV2 Class with Custom Configuration (Python)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

Shows how to create a ConvNeXtV2 model instance using the main class with custom configurations for input channels, number of classes, stage depths, dimensions, and drop path rate. Includes examples for training and inference.

```python
import torch
from models.convnextv2 import ConvNeXtV2

# Custom configuration for ConvNeXt V2
model = ConvNeXtV2(
    in_chans=3,
    num_classes=1000,
    depths=[3, 3, 27, 3],
    dims=[128, 256, 512, 1024],
    drop_path_rate=0.1,
    head_init_scale=1.0
)

# Model structure:
# - Stem: 4x4 conv with stride 4
# - 4 stages with downsampling between them
# - Global average pooling + linear classifier

# Training example
model.train()
images = torch.randn(8, 3, 224, 224)
targets = torch.randint(0, 1000, (8,))
criterion = torch.nn.CrossEntropyLoss()

outputs = model(images)
loss = criterion(outputs, targets)
loss.backward()

# Inference example
model.eval()
with torch.no_grad():
    predictions = model(images).argmax(dim=1)
    print(f"Predicted classes: {predictions}")
```

--------------------------------

### Evaluate ConvNeXt V2 Model (Single-GPU)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/README.md

This command demonstrates how to evaluate a ConvNeXt V2 model on a single GPU. It requires specifying the model variant, evaluation mode, path to the checkpoint, input size, and dataset path. The `main_finetune.py` script is used for this purpose.

```python
python main_finetune.py \
--model convnextv2_base \
--eval true \
--resume /path/to/checkpoint \
--input_size 224 \
--data_path /path/to/imagenet-1k \

```

--------------------------------

### Evaluate ConvNeXt V2 Model (Multi-GPU)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/README.md

This command shows how to evaluate a ConvNeXt V2 model using multiple GPUs. It utilizes `torch.distributed.launch` to manage distributed training and requires similar parameters as the single-GPU evaluation, including model, checkpoint, input size, and data path.

```python
python -m torch.distributed.launch --nproc_per_node=8 main_finetune.py \
--model convnextv2_base \
--eval true \
--resume /path/to/checkpoint \
--input_size 224 \
--data_path /path/to/imagenet-1k \

```

--------------------------------

### FCMAE Pre-Training on Multi-Node SLURM (Bash)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

This command demonstrates distributed FCMAE pre-training across multiple nodes using the SLURM job scheduler and submitit. It specifies the number of nodes, GPUs per node, and essential training configurations.

```bash
# Pre-train ConvNeXt V2 Base on ImageNet-1K (8 nodes x 8 GPUs)
python submitit_pretrain.py \
    --nodes 8 \
    --ngpus 8 \
    --model convnextv2_base \
    --batch_size 64 \
    --blr 1.5e-4 \
    --epochs 1600 \
    --warmup_epochs 40 \
    --data_path /path/to/imagenet-1k \
    --job_dir /path/to/save_results
```

--------------------------------

### Build Datasets and Transforms (Python)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

This snippet shows how to import functions for building datasets and applying transformations, likely for preparing data for model training or evaluation. It imports `build_dataset` and `build_transform` from the `datasets` module and `Namespace` from `argparse`.

```python
from datasets import build_dataset, build_transform
from argparse import Namespace
```

--------------------------------

### FCMAE Pre-Training on Single Machine (Multi-GPU, Bash)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

This command initiates FCMAE pre-training for ConvNeXt V2 Base on ImageNet-1K using multiple GPUs on a single machine. It configures training parameters such as batch size, learning rate, epochs, and masking ratio.

```bash
# Pre-train ConvNeXt V2 Base on ImageNet-1K (single machine, 8 GPUs)
python -m torch.distributed.launch --nproc_per_node=8 main_pretrain.py \
    --model convnextv2_base \
    --batch_size 64 \
    --update_freq 8 \
    --blr 1.5e-4 \
    --epochs 1600 \
    --warmup_epochs 40 \
    --mask_ratio 0.6 \
    --decoder_depth 1 \
    --decoder_embed_dim 512 \
    --norm_pix_loss \
    --data_path /path/to/imagenet-1k \
    --output_dir /path/to/save_results
```

--------------------------------

### Create and Use ConvNeXt V2 Model Variants (Python)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

Demonstrates how to instantiate different ConvNeXt V2 model sizes (e.g., Atto, Femto, Base) using factory functions. It shows how to perform a forward pass for classification and feature extraction.

```python
import torch
from models.convnextv2 import (
    convnextv2_atto,
    convnextv2_femto,
    convnext_pico,
    convnextv2_nano,
    convnextv2_tiny,
    convnextv2_base,
    convnextv2_large,
    convnextv2_huge,
)

# Create a ConvNeXt V2 Base model for 1000-class classification
model = convnextv2_base(
    num_classes=1000,
    drop_path_rate=0.1,
    head_init_scale=0.001
)

# Forward pass with batch of images
images = torch.randn(4, 3, 224, 224)  # (batch, channels, height, width)
logits = model(images)  # Output shape: (4, 1000)

# Extract features without classification head
features = model.forward_features(images)  # Output shape: (4, 1024)
```

--------------------------------

### Pre-train ConvNeXt V2-Base on ImageNet-1K (Multi-node)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Command to pre-train the ConvNeXt V2-Base model on ImageNet-1K using 8 nodes, each with 8 GPUs. It specifies training parameters like batch size, learning rate, epochs, and data/output paths. Requires submitit for execution.

```bash
python submitit_pretrain.py --nodes 8 --ngpus 8 \
--model convnextv2_base \
--batch_size 64 \
--blr 1.5e-4 \
--epochs 1600 \
--warmup_epochs 40 \
--data_path /path/to/imagenet-1k \
--job_dir /path/to/save_results
```

--------------------------------

### Implement FCMAE Model for Pre-training (Python)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

This snippet demonstrates how to instantiate and use the FCMAE class for self-supervised pre-training. It defines model parameters, performs a forward pass with random images, and explains the output tensors (loss, predictions, and mask).

```python
import torch
from models.fcmae import FCMAE

# Create FCMAE model for pre-training
fcmae = FCMAE(
    img_size=224,
    in_chans=3,
    depths=[3, 3, 27, 3],      # Encoder depths (Base config)
    dims=[128, 256, 512, 1024], # Encoder dimensions
    decoder_depth=1,            # Number of decoder blocks
    decoder_embed_dim=512,      # Decoder embedding dimension
    patch_size=32,              # Patch size for masking
    mask_ratio=0.6,             # Ratio of patches to mask
    norm_pix_loss=True          # Normalize pixel values for loss
)

# Pre-training forward pass
images = torch.randn(4, 3, 224, 224)
loss, pred, mask = fcmae(images, mask_ratio=0.6)

# Output explanation:
# - loss: Reconstruction loss on masked patches only
# - pred: Predicted pixel values shape (4, num_patches, patch_size^2 * 3)
# - mask: Binary mask (4, num_patches), 1 = masked, 0 = visible

print(f"Pre-training loss: {loss.item():.4f}")
print(f"Prediction shape: {pred.shape}")
print(f"Mask shape: {mask.shape}")
```

--------------------------------

### Load Pre-Trained Weights for ConvNeXt V2 Models

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

Demonstrates loading pre-trained weights for ConvNeXt V2 models. It covers loading fine-tuned checkpoints, handling potential shape mismatches in the classification head, and loading FCMAE pre-trained weights by removing decoder components and remapping keys.

```python
import torch
from utils import remap_checkpoint_keys, load_state_dict
import models.convnextv2 as convnextv2

# Create model
model = convnextv2.convnextv2_base(
    num_classes=1000,
    drop_path_rate=0.1,
    head_init_scale=0.001
)

# Load fine-tuned checkpoint
checkpoint = torch.load('/path/to/convnextv2_base_1k_224_ema.pt', map_location='cpu')
checkpoint_model = checkpoint['model']

# Remove head weights if shape mismatch (e.g., different number of classes)
state_dict = model.state_dict()
for k in ['head.weight', 'head.bias']:
    if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
        del checkpoint_model[k]

# Load state dict
load_state_dict(model, checkpoint_model)
model.eval()

# For FCMAE pre-trained weights, remove decoder components
checkpoint = torch.load('/path/to/fcmae_pretrained.pt', map_location='cpu')
checkpoint_model = checkpoint['model']

# Remove decoder-specific keys
for k in list(checkpoint_model.keys()):
    if 'decoder' in k or 'mask_token' in k or 'proj' in k or 'pred' in k:
        del checkpoint_model[k]

# Remap sparse encoder keys to dense model keys
checkpoint_model = remap_checkpoint_keys(checkpoint_model)
load_state_dict(model, checkpoint_model)
```

--------------------------------

### Pre-train ConvNeXt V2-Base on ImageNet-1K (Single Machine)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Command to pre-train the ConvNeXt V2-Base model on ImageNet-1K on a single machine with 8 GPUs. It utilizes torch.distributed.launch for distributed training and specifies various training hyperparameters.

```bash
python -m torch.distributed.launch --nproc_per_node=8 main_pretrain.py \
--model convnextv2_base \
--batch_size 64 --update_freq 8 \
--blr 1.5e-4 \
--epochs 1600 \
--warmup_epochs 40 \
--data_path /path/to/imagenet-1k \
--output_dir /path/to/save_results
```

--------------------------------

### Fine-tune ConvNeXt V2-Atto on ImageNet-1K (Single Machine)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Command for fine-tuning the ConvNeXt V2-Atto model on ImageNet-1K on a single machine with 8 GPUs. It utilizes torch.distributed.launch and sets specific hyperparameters for the Atto model, such as layer decay and augmentation settings.

```bash
python -m torch.distributed.launch --nproc_per_node=8 main_finetune.py \
--model convnextv2_atto \
--batch_size 32 --update_freq 4 \
--blr 2e-4 \
--epochs 600 \
--warmup_epochs 0 \
--layer_decay_type 'single' \
--layer_decay 0.9 \
--weight_decay 0.3 \
--drop_path 0.1 \
--reprob 0.25 \
--mixup 0. \
--cutmix 0. \
--smoothing 0.2 \
--model_ema True --model_ema_eval True \
--use_amp True \
--finetune /path/to/checkpoint \
--data_path /path/to/imagenet-1k \
--output_dir /path/to/save_results
```

--------------------------------

### JAX Block Module for FCMAE Pre-training

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Implements the Block module for FCMAE pre-training in JAX. This module utilizes DepthwiseConv2D, LayerNorm, Dense layers, GELU activation, and the GRN module. It also includes dropout for regularization and supports optional masking, which is noted to be numerically equivalent to sparse convolution.

```python
from flax import linen as nn
import jax.numpy as jnp

class Block(nn.Module):
  dim: int
  drop_path: float
  
  @nn.compact
  def __call__(self, inputs, mask=None):
    if mask is not None:
      x = inputs * (1. - mask)
    x = DepthwiseConv2D((7, 7), name='dwconv')(x)
    if mask is not None: # The binary masking is numerically identical to sparse conv. 
      x = x * (1.- mask)
    x = nn.LayerNorm(name='norm')(x)
    x = nn.Dense(4 * self.dim, name='pwconv1')(x)
    x = nn.gelu(x)
    x = GRN(4 * self.dim, name='grn')(x, mask)
    x = nn.Dense(self.dim, name='pwconv2')(x)
    x = nn.Dropout(rate=self.drop_path, broadcast_dims=(1,2,3), name='droppath')(x, deterministic=not self.training)
    return x + inputs
```

--------------------------------

### Fine-Tuning ConvNeXt V2 on Single Machine (Multi-GPU, Bash)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

This command fine-tunes a pre-trained FCMAE checkpoint for image classification on a single machine with multiple GPUs. It includes various data augmentation, regularization, and optimization techniques.

```bash
# Fine-tune ConvNeXt V2 Base on ImageNet-1K (single machine, 8 GPUs)
python -m torch.distributed.launch --nproc_per_node=8 main_finetune.py \
    --model convnextv2_base \
    --batch_size 32 \
    --update_freq 4 \
    --blr 6.25e-4 \
    --epochs 100 \
    --warmup_epochs 20 \
    --layer_decay_type group \
    --layer_decay 0.6 \
    --weight_decay 0.05 \
    --drop_path 0.1 \
    --reprob 0.25 \
    --mixup 0.8 \
    --cutmix 1.0 \
    --smoothing 0.1 \
    --model_ema true \
    --model_ema_eval true \
    --use_amp true \
    --finetune /path/to/fcmae_checkpoint.pt \
    --data_path /path/to/imagenet-1k \
    --output_dir /path/to/save_results
```

--------------------------------

### Fine-tune ConvNeXt V2-Base on ImageNet-1K (Multi-node)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Command to fine-tune the ConvNeXt V2-Base model on ImageNet-1K using 4 nodes, each with 8 GPUs. It includes parameters for fine-tuning such as batch size, learning rate, epochs, weight decay, and specifies a pre-trained checkpoint.

```bash
python submitit_finetune.py --nodes 4 --ngpus 8 \
--model convnextv2_base \
--batch_size 32 \
--blr 6.25e-4 \
--epochs 100 \
--warmup_epochs 20 \
--layer_decay_type 'group' \
--layer_decay 0.6 \
--weight_decay 0.05 \
--drop_path 0.1 \
--reprob 0.25 \
--mixup 0.8 \
--cutmix 1.0 \
--smoothing 0.1 \
--model_ema True --model_ema_eval True \
--use_amp True \
--finetune /path/to/checkpoint \
--data_path /path/to/imagenet-1k \
--job_dir /path/to/save_results
```

--------------------------------

### Fine-tune ConvNeXt V2-Tiny on ImageNet-1K (Single Machine)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Command for fine-tuning the ConvNeXt V2-Tiny model on ImageNet-1K on a single machine with 8 GPUs. It uses torch.distributed.launch and sets hyperparameters specific to the Tiny model, including learning rate and augmentation parameters.

```bash
python -m torch.distributed.launch --nproc_per_node=8 main_finetune.py \
--model convnextv2_ \
--batch_size 32 --update_freq 4 \
--blr 8e-4 \
--epochs 300 \
--warmup_epochs 40 \
--layer_decay_type 'single' \
--layer_decay 0.9 \
--weight_decay 0.05 \
--drop_path 0.2 \
--reprob 0.25 \
--mixup 0.8 \
--cutmix 1.0 \
--smoothing 0.1 \
--model_ema True --model_ema_eval True \
--use_amp True \
--finetune /path/to/checkpoint \
--data_path /path/to/imagenet-1k \
--output_dir /path/to/save_results
```

--------------------------------

### Evaluate ConvNeXt V2 Model on ImageNet-1K (Single GPU, Bash)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

This command demonstrates how to evaluate a fine-tuned ConvNeXt V2 model on the ImageNet-1K validation set using a single GPU. It specifies the model, evaluation mode, path to the pre-trained checkpoint, input size, and dataset path.

```bash
# Single-GPU evaluation
python main_finetune.py \
    --model convnextv2_base \
    --eval true \
    --resume /path/to/convnextv2_base_1k_224_ema.pt \
    --input_size 224 \
    --data_path /path/to/imagenet-1k

# Expected output:
# * Acc@1 84.900 Acc@5 97.250 loss 0.678
```

--------------------------------

### Fine-tune ConvNeXt V2-Base on ImageNet-1K (Single Machine)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Command to fine-tune the ConvNeXt V2-Base model on ImageNet-1K on a single machine with 8 GPUs. It uses torch.distributed.launch and specifies fine-tuning hyperparameters and the path to the pre-trained model.

```bash
python -m torch.distributed.launch --nproc_per_node=8 main_finetune.py \
--model convnextv2_base \
--batch_size 32 --update_freq 4 \
--blr 6.25e-4 \
--epochs 100 \
--warmup_epochs 20 \
--layer_decay_type 'group' \
--layer_decay 0.6 \
--weight_decay 0.05 \
--drop_path 0.1 \
--reprob 0.25 \
--mixup 0.8 \
--cutmix 1.0 \
--smoothing 0.1 \
--model_ema True --model_ema_eval True \
--use_amp True \
--finetune /path/to/checkpoint \
--data_path /path/to/imagenet-1k \
--output_dir /path/to/save_results
```

--------------------------------

### Configure and Build Datasets for ConvNeXt V2

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

Configures dataset parameters including data path, augmentation policies (AutoAugment, RandAugment), random erasing, and interpolation. It then builds training and validation datasets using the `build_dataset` function and separate transforms using `build_transform`.

```python
from argparse import Namespace

# Configuration for dataset building
args = Namespace(
    data_path='/path/to/imagenet-1k',
    data_set='IMNET',
    input_size=224,
    color_jitter=0.4,
    aa='rand-m9-mstd0.5-inc1',  # AutoAugment policy
    train_interpolation='bicubic',
    reprob=0.25,   # Random erasing probability
    remode='pixel',
    recount=1,
    imagenet_default_mean_and_std=True,
    crop_pct=None,
    eval_data_path=None,
    nb_classes=1000
)

# Build training dataset with augmentation
train_dataset, num_classes = build_dataset(is_train=True, args=args)
print(f"Training samples: {len(train_dataset)}, Classes: {num_classes}")

# Build validation dataset
val_dataset, _ = build_dataset(is_train=False, args=args)
print(f"Validation samples: {len(val_dataset)}")

# Build transforms separately
train_transform = build_transform(is_train=True, args=args)
val_transform = build_transform(is_train=False, args=args)
```

--------------------------------

### Fine-Tuning ConvNeXt V2 on Multi-Node SLURM (Bash)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

This command performs large-scale fine-tuning of ConvNeXt V2 across multiple nodes using SLURM. It configures distributed training parameters and essential fine-tuning settings for efficient training on large datasets.

```bash
# Fine-tune ConvNeXt V2 Base on ImageNet-1K (4 nodes x 8 GPUs)
python submitit_finetune.py \
    --nodes 4 \
    --ngpus 8 \
    --model convnextv2_base \
    --batch_size 32 \
    --blr 6.25e-4 \
    --epochs 100 \
    --warmup_epochs 20 \
    --layer_decay_type group \
    --layer_decay 0.6 \
    --weight_decay 0.05 \
    --drop_path 0.1 \
    --reprob 0.25 \
    --mixup 0.8 \
    --cutmix 1.0 \
    --smoothing 0.1 \
    --model_ema true \
    --model_ema_eval true \
    --use_amp true \
    --finetune /path/to/fcmae_checkpoint.pt \
    --data_path /path/to/imagenet-1k \
    --job_dir /path/to/save_results
```

--------------------------------

### Evaluate ConvNeXt V2 Model on ImageNet-1K (Multi-GPU, Bash)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

This command shows how to perform distributed evaluation of a ConvNeXt V2 model across multiple GPUs. It utilizes `torch.distributed.launch` to manage the distributed processes and enables distributed evaluation with `--dist_eval true`.

```bash
# Multi-GPU evaluation with 8 GPUs
python -m torch.distributed.launch --nproc_per_node=8 main_finetune.py \
    --model convnextv2_base \
    --eval true \
    --resume /path/to/convnextv2_base_1k_224_ema.pt \
    --input_size 224 \
    --data_path /path/to/imagenet-1k \
    --dist_eval true
```

--------------------------------

### ConvNeXt V2 Training and Evaluation Functions

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

Sets up training and evaluation pipelines using provided functions. It configures Mixup augmentation and SoftTargetCrossEntropy loss for training, then calls `train_one_epoch` and `evaluate` to perform one training epoch and model evaluation, respectively.

```python
import torch
from engine_finetune import train_one_epoch, evaluate
from timm.data import Mixup
from timm.loss import SoftTargetCrossEntropy

# Setup mixup augmentation
mixup_fn = Mixup(
    mixup_alpha=0.8,
    cutmix_alpha=1.0,
    prob=1.0,
    switch_prob=0.5,
    mode='batch',
    label_smoothing=0.1,
    num_classes=1000
)

# Training criterion with mixup
criterion = SoftTargetCrossEntropy()

# Train one epoch
train_stats = train_one_epoch(
    model=model,
    criterion=criterion,
    data_loader=train_loader,
    optimizer=optimizer,
    device=torch.device('cuda'),
    epoch=0,
    loss_scaler=loss_scaler,
    max_norm=None,  # Gradient clipping
    model_ema=model_ema,
    mixup_fn=mixup_fn,
    log_writer=None,
    args=args
)
print(f"Training loss: {train_stats['loss']:.4f}")

# Evaluate model
test_stats = evaluate(
    data_loader=val_loader,
    model=model,
    device=torch.device('cuda'),
    use_amp=True
)
print(f"Top-1 Accuracy: {test_stats['acc1']:.2f}%")
print(f"Top-5 Accuracy: {test_stats['acc5']:.2f}%")
```

--------------------------------

### Fine-tune ConvNeXt V2-Atto on ImageNet-1K (Multi-node)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Command for fine-tuning the ConvNeXt V2-Atto model on ImageNet-1K across 4 nodes with 8 GPUs each. It configures specific hyperparameters for the Atto model, including layer decay and mixup/cutmix settings.

```bash
python submitit_finetune.py --nodes 4 --ngpus 8 \
--model convnextv2_atto \
--batch_size 32 \
--blr 2e-4 \
--epochs 600 \
--warmup_epochs 0 \
--layer_decay_type 'single' \
--layer_decay 0.9 \
--weight_decay 0.3 \
--drop_path 0.1 \
--reprob 0.25 \
--mixup 0. \
--cutmix 0. \
--smoothing 0.2 \
--model_ema True --model_ema_eval True \
--use_amp True \
--finetune /path/to/checkpoint \
--data_path /path/to/imagenet-1k \
--job_dir /path/to/save_results
```

--------------------------------

### Fine-tune ConvNeXt V2-Tiny on ImageNet-1K (Multi-node)

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Command for fine-tuning the ConvNeXt V2-Tiny model on ImageNet-1K across 4 nodes with 8 GPUs each. It configures hyperparameters tailored for the Tiny model, including learning rate, epochs, and augmentation strategies.

```bash
python submitit_finetune.py --nodes 4 --ngpus 8 \
--model convnextv2_tiny \
--batch_size 32 \
--blr 8e-4 \
--epochs 300 \
--warmup_epochs 40 \
--layer_decay_type 'single' \
--layer_decay 0.9 \
--weight_decay 0.05 \
--drop_path 0.2 \
--reprob 0.25 \
--mixup 0.8 \
--cutmix 1.0 \
--smoothing 0.1 \
--model_ema True --model_ema_eval True \
--use_amp True \
--finetune /path/to/checkpoint \
--data_path /path/to/imagenet-1k \
--job_dir /path/to/save_results
```

--------------------------------

### JAX GRN Module for FCMAE Pre-training

Source: https://github.com/facebookresearch/convnext-v2/blob/main/TRAINING.md

Defines the Gated Residual Network (GRN) module in JAX, used within the FCMAE pre-training framework. This module incorporates learnable parameters (gamma and beta) and applies normalization and gating mechanisms. It handles optional masking for specific layers.

```python
from flax import linen as nn
import jax.numpy as jnp

class GRN(nn.Module):
  dim: int
  eps: float = 1e-6
  
  def init_fn(self, key, shape, fill_value):
    return jnp.full(shape, fill_value)
  
  @nn.compact
  def __call__(self, inputs, mask=None):
    gamma = self.param("gamma", self.init_fn, (self.dim,), 0.)
    beta = self.param("beta", self.init_fn, (self.dim,), 0.)
    
    x = inputs
    if mask is not None:
      x = x * (1. - mask)
    GX = jnp.power((jnp.sum(jnp.power(x, 2), axis=(1,2), keepdims=True) + self.eps), 0.5)
    Nx = Gx / (jnp.mean(Gx, axis=-1, keepdims=True) + self.eps)
    return gamma * (Nx * inputs) + beta + inputs
```

--------------------------------

### Utilize GRN Layer for Inter-Channel Feature Competition (Python)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

Demonstrates the usage of the Global Response Normalization (GRN) layer, a key component of ConvNeXt V2. It explains how GRN enhances feature competition by normalizing channel responses based on global statistics and shows its learnable parameters.

```python
import torch
from models.utils import GRN

# Create GRN layer for 256 channels
grn = GRN(dim=256)

# GRN expects input shape: (batch, height, width, channels)
# This is the channels_last format used after permutation in ConvNeXt blocks
x = torch.randn(4, 56, 56, 256)

# Forward pass
# GRN computes: gamma * (x * Nx) + beta + x
# where Nx = Gx / mean(Gx), Gx = L2_norm(x) across spatial dims
output = grn(x)  # Shape: (4, 56, 56, 256)

# GRN has learnable parameters gamma and beta
print(f"Gamma shape: {grn.gamma.shape}")
print(f"Beta shape: {grn.beta.shape}")
```

--------------------------------

### Custom LayerNorm Implementation for ConvNeXt (Python)

Source: https://context7.com/facebookresearch/convnext-v2/llms.txt

Details the custom LayerNorm implementation within ConvNeXt V2, which supports both 'channels_last' and 'channels_first' data formats. This flexibility is crucial for different parts of the network, such as the stem and internal blocks.

```python
import torch
from models.utils import LayerNorm

# Channels-last format (default) - used in block internals
ln_last = LayerNorm(normalized_shape=256, eps=1e-6, data_format="channels_last")
x_last = torch.randn(4, 56, 56, 256)  # (batch, H, W, C)
output_last = ln_last(x_last)  # Shape: (4, 56, 56, 256)

# Channels-first format - used in stem and downsampling layers
ln_first = LayerNorm(normalized_shape=256, eps=1e-6, data_format="channels_first")
x_first = torch.randn(4, 256, 56, 56)  # (batch, C, H, W)
output_first = ln_first(x_first)  # Shape: (4, 256, 56, 56)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.