TR-HASH Vision V6 · ImageNet-1K Pretrain

This is the completed ImageNet-1K pretraining release for the compact TR-HASH Vision V6 tower. Its primary purpose is to initialize downstream V6 detectors. It also includes the classification head used for the reported ImageNet-1K validation result.

This repository is a vision backbone pretrain, not a finished object detector. Detection checkpoints and mAP results are published separately.

Result
Full classifier 0.764M parameters
Reusable vision tower 0.635M parameters
Input resolution 224 × 224
Training 100 epochs / 500,500 steps
ImageNet-1K top-1 accuracy 50.834%
Training hardware 4× RTX 5060 Ti 16 GB
Wall time 11 h 01 min

Training curves

TR-HASH Vision V6 ImageNet-1K training curves

The loss curve uses all 10,026 logged training observations, with a 100-log moving average. Validation was run every five epochs on all 50,000 ImageNet-1K validation images. The complete raw history is available in metrics.jsonl.

Epoch Step Top-1 accuracy
5 25,025 32.152%
20 100,100 42.904%
40 200,200 45.212%
60 300,300 47.888%
80 400,400 49.396%
100 500,500 50.834%

Architecture

TR-HASH Vision V6 is a compact hierarchical vision transformer whose spatial token IDs are hashed into expert selections at each routed layer. Each token activates two of four experts while shared attention continues to mix the full contextual feature map.

  • Three native feature stages with depths 1 / 1 / 2
  • Patch size 8, hidden width 128, four attention heads
  • Shifted local windows of size 8
  • Four routed experts, top-2 activation, expert width 48
  • BF16 training with fused CUDA/Triton routed kernels
  • Shared tower layout used directly by the V6 detection architecture
spatial token ID ──► layer-specific hash routing ──► selected expert weights
        │                                                   │
        └────────── contextual attention state ─────────────┴──► output

Files

  • tower.safetensors — reusable 0.635M-parameter tower for detection and feature extraction.
  • model.safetensors — 0.764M-parameter ImageNet-1K classifier, including normalization and the 1,000-class head.
  • config.json — exact V6 tower configuration and validated epoch metadata.
  • class_names.json — the 1,000 standard ILSVRC class names in output order.
  • preprocessor_config.json — validation preprocessing parameters.
  • metrics.jsonl — raw loss, learning-rate and validation-accuracy history.
  • training_summary.json — compact machine-readable training summary.
  • training_state.pt — optimizer, scheduler, epoch/step cursor and RNG state for an exact framework resume. This PyTorch state file is not required for inference; load it only from this trusted repository.

Load the classifier

import json

import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from torchvision import transforms

from complexity.generative.detection.config import TRHashDetectorConfig
from complexity.generative.detection.hierarchical_tower import (
    HierarchicalTRHashVisionClassifier,
)

repo = snapshot_download(
    "AETHORIA-AI/TR-HASH-Vision-V6-ImageNet1K-Pretrain"
)
metadata = json.load(open(f"{repo}/config.json"))
config = TRHashDetectorConfig.from_dict(metadata["tower"])

model = HierarchicalTRHashVisionClassifier(config, num_classes=1000).eval()
model.load_state_dict(load_file(f"{repo}/model.safetensors"))

preprocess = transforms.Compose([
    transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

with torch.inference_mode():
    logits = model(preprocess(image).unsqueeze(0))["logits"]
    predicted_class = int(logits.argmax(dim=-1))

Install the framework from the source repository:

pip install "git+https://github.com/Complexity-ML/complexity-framework.git"

Initialize a V6 detector

The detection trainer consumes this repository's tower weights through --backbone-checkpoint after downloading the snapshot locally:

python -m complexity.generative.detection.training \
  --backbone-checkpoint /path/to/TR-HASH-Vision-V6-ImageNet1K-Pretrain \
  --architecture-version 6 \
  --image-size 640 \
  --patch-size 8 \
  --vision-hidden-size 128 \
  --vision-stage-depths 1 1 2 \
  --vision-window-size 8 \
  --vision-heads 4 \
  --vision-num-experts 4 \
  --vision-top-k 2 \
  --vision-expert-width 48 \
  ...

Resume pretraining

training_state.pt, model.safetensors and config.json form the complete epoch-100 training checkpoint. To extend the run with the same recipe, download the repository and point the V6 pretraining launcher at its directory:

RESUME_CHECKPOINT=/path/to/TR-HASH-Vision-V6-ImageNet1K-Pretrain \
  scripts/vast_pretrain_vision_v06_imagenet1k.sh

Changing the world size, model configuration or training recipe is not an exact resume and should instead be treated as a new fine-tuning run.

Training recipe

  • Dataset: ImageNet-1K, using the public 256 px Hugging Face repack benjamin-paine/imagenet-1k-256x256
  • Samples: 1,281,167 train and 50,000 validation images
  • Optimizer: fused AdamW, weight decay 0.05
  • LR: 3e-4 base and 4.5e-4 routed experts
  • Schedule: 5,000-step warmup followed by cosine decay
  • Batch: 64 images per GPU, global batch 256
  • Augmentation: random resized crop, horizontal flip and color jitter
  • Precision/runtime: BF16, PyTorch DDP, NCCL and fused CUDA/Triton kernels
  • Seed: 3

Intended use and limitations

This release is intended for research, architecture ablations, compact feature extraction and initialization of TR-HASH Vision V6 downstream models. Its 50.834% ImageNet-1K top-1 result is a realized baseline for this sub-million parameter architecture; it is not a state-of-the-art classification claim and does not establish detection quality.

The model may inherit biases and failure modes from ImageNet. Validate accuracy, calibration, robustness and downstream behavior on your target data before use. ImageNet images are not redistributed here and remain subject to their original terms.

Links

License

The released model artifacts and repository materials are provided under CC BY-NC 4.0.

Downloads last month
-
Safetensors
Model size
767k params
Tensor type
I64
·
F32
·
I32
·
U8
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train AETHORIA-AI/TR-HASH-Vision-V6-ImageNet1K-Pretrain

Collection including AETHORIA-AI/TR-HASH-Vision-V6-ImageNet1K-Pretrain

Evaluation results