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
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 width128, 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-4base and4.5e-4routed 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
- -
Dataset used to train AETHORIA-AI/TR-HASH-Vision-V6-ImageNet1K-Pretrain
Collection including AETHORIA-AI/TR-HASH-Vision-V6-ImageNet1K-Pretrain
Evaluation results
- Top-1 accuracy on ImageNet-1K validationself-reported0.508
