Custom Vision Transformer Workload

This repository contains an original, from-scratch implementation of a reusable Vision Transformer architecture in PyTorch. Its purpose is to provide a controlled transformer workload for mapping primitive data movement and concurrency across the operations that compose modern transformer and LLM systems.

The architecture moves the broader benchmarking project above isolated low-level primitives and into a complete model execution graph. Patch projection, attention, normalization, feed-forward computation, residual paths, and prediction heads create a practical layer in which memory movement, parallelism, synchronization, and hardware utilization can be related back to model-level behavior.

The current proof of concept has been trained and tested on a sample of the Railway Fastener Defect Dataset. In this application, the model predicts defect classes and normalized bounding boxes for railway fastener components.

The transformer itself is not railway-specific. Its input pipeline, prediction head, loss, and dataset adapter can be changed to reuse the encoder for other computer-vision applications.

Project context

This transformer was built to complete the architectural portion of Phase 1 of the broader GPU Benchmarking project, expanding Phase 0's primitive-level work into the higher abstraction layer used by transformer and LLM workloads.

The custom implementation keeps important operations visible and modifiable, making it suitable for tracing how tensors move through the model and where concurrent execution is created, limited, or synchronized. The detailed data movement and concurrency mapping is planned work; this repository currently establishes the reusable architecture, training workload, inference path, and visual outputs needed to perform that analysis.

Architecture

Stage	Component	Input	Output
1	Image patching	[B, C, H, W]	[B, N, patch_dim]
2	Patch projection, positional embedding, and CLS token	[B, N, patch_dim]	[B, N+1, d_model]
3	Custom transformer encoder stack	[B, N+1, d_model]	[B, N+1, d_model]
4	Spatial pooling and multi-object head	[B, N, d_model]	[B, object_slots, 4 + classes + background]
5	Multitask training	Predictions and targets	Classification and box-regression loss

The implementation includes custom multi-head attention, layer normalization, feed-forward blocks, residual connections, patch extraction, and a detection head. Keeping these components explicit—rather than hiding the full workload behind a prebuilt model—supports future operation-level instrumentation and concurrency analysis. Model and dataset settings are centralized in Transformer/src/config.py.

Inference visualization

The inference workflow overlays normalized model predictions and ground-truth annotations on validation images. Ground-truth boxes are shown in green and predictions in red.

These images demonstrate the complete data-to-visualization pipeline. They are experimental outputs rather than a claim of production-level accuracy.

Running the project

From the Transformer directory, train the model with:

PYTHONPATH=. python -m src.run

Run validation inference and generate visualization images with:

PYTHONPATH=. python -m pred.inference

Generated overlays are saved under Transformer/pred/visualizations/.