VisionForge

VisionForge is a production-oriented PyTorch platform for explainable industrial visual inspection. It is organized as reusable software rather than a notebook: data manifests, model paths, training, evaluation, checkpoints, APIs, local UI, and model publishing all live in one repository.

Model paths

• scratch_cnn: transparent CNN baseline trained from random initialization.
• resnet18 and vit: pretrained transfer-learning baselines.
• autoencoder: one-class reconstruction anomaly detector, trained only on defect-free samples.

The intended dataset is MVTec AD. Metrics include classification accuracy/precision/recall/F1/ROC-AUC and anomaly ROC-AUC/AP, with inference latency returned by every prediction.

Install

cd ~/Documents/visionforge
python -m venv .venv
source .venv/bin/activate
pip install -e '.[dev]'

Download and extract MVTec AD so data/raw/mvtec_ad/bottle/train/good exists. Create auditable data manifests:

python scripts/prepare_mvtec.py --root data/raw/mvtec_ad --category bottle

Train

# Scratch CNN, one-epoch smoke run
python -m visionforge.training.train_classification --config configs/mvtec/smoke_classification.yaml

# ResNet-18 transfer learning
python -m visionforge.training.train_classification --config configs/mvtec/resnet18.yaml

# Strict one-class reconstruction anomaly baseline
python -m visionforge.training.train_anomaly --config configs/mvtec/autoencoder.yaml

The classification entry point is a supervised benchmark: it uses explicitly labelled MVTec images and random train/validation partitioning. It must not be reported as an unseen-defect protocol. Use train_anomaly for the standard setting where only defect-free training images are available.

Serve locally

VISIONFORGE_CHECKPOINT=outputs/smoke_classification/best.pt \
  uvicorn visionforge.serving.api:app --host 0.0.0.0 --port 8000

VISIONFORGE_CHECKPOINT=outputs/smoke_classification/best.pt \
  python -m visionforge.serving.gradio_app

GET /health reports model readiness. POST /predict accepts a multipart image and returns prediction, confidence, anomaly_score, and latency_ms. The Gradio app displays the same result with a Grad-CAM heatmap for convolutional classifiers.

Publish a model

python scripts/export_hf.py outputs/smoke_classification/best.pt --repo-id your-user/visionforge-bottle

The repository intentionally keeps datasets, checkpoints, and generated outputs out of version control. Before deployment, calibrate the anomaly threshold on representative normal products and test across camera, lighting, material, and defect variations.

Distributed AWS training

VisionForge uses the same Accelerate-style rank-safe training pattern as AudioForge: every worker trains a shard, while only rank zero writes checkpoints, metrics, and runtime metadata. batch_size is per GPU, so global batch size is batch_size × GPU count × gradient_accumulation_steps.

For one AWS GPU host with two or more GPUs, copy the repository and MVTec data to the instance, install with pip install -e '.[dev]', then run:

VISIONFORGE_DATA_ROOT=/mnt/mvtec_ad \
VISIONFORGE_CONFIG=configs/mvtec/resnet18_aws.yaml \
VISIONFORGE_NUM_PROCESSES=2 \
bash scripts/train_2gpu.sh

For multi-node EC2, place the checkout and output directory on a shared filesystem (FSx for Lustre/EFS) or configure S3 sync. Open TCP MASTER_PORT only within the security group, and run on every node with its correct rank:

NNODES=2 NODE_RANK=0 MASTER_ADDR=10.0.1.10 NPROC_PER_NODE=4 \
VISIONFORGE_DATA_ROOT=/fsx/mvtec_ad VISIONFORGE_OUTPUT_DIR=/fsx/visionforge-output \
bash scripts/aws/train_ec2_distributed.sh

On the other node, use NODE_RANK=1 and the same remaining values. To copy rank-zero artifacts to S3 when the job finishes, set VISIONFORGE_SYNC_S3_URI=s3://bucket/visionforge/run-1.

For managed distributed training, install the optional SDK locally and submit a job. The S3 training prefix must contain category folders directly (for example bottle/train/good).

The submitted source bundle includes requirements.txt; SageMaker installs these project-level runtime packages alongside its managed PyTorch image.

pip install -e '.[aws]'
python scripts/aws/submit_sagemaker.py \
  --role arn:aws:iam::ACCOUNT_ID:role/SageMakerExecutionRole \
  --data-s3-uri s3://YOUR_BUCKET/mvtec_ad/ \
  --output-s3-uri s3://YOUR_BUCKET/visionforge-output/ \
  --instance-type ml.g5.2xlarge --instance-count 2

The training container recipe is docker/Dockerfile.train. Use it for EC2/ECS-compatible builds when you need an immutable CUDA runtime.