feat(batch27): JEPA ViT bridge, WebDataset shard reading, training loop
CI / Format Check (push) Failing after 11s
CI / Test (macos-latest) (push) Has been skipped
CI / Test (ubuntu-latest) (push) Has been skipped
CI / Python Bindings (maturin) (macos-latest) (push) Has been skipped
CI / Python Bindings (maturin) (ubuntu-latest) (push) Has been skipped
CI / WASM Build + Size Check (push) Has been skipped
Documentation / Build User Guide (push) Successful in 9s
CI / Build CPU-Only (Explicit) (push) Failing after 1m10s
CI / CI Success (push) Failing after 0s
Documentation / Build API Documentation (push) Failing after 40s
Performance Benchmarks / Run Benchmarks (push) Successful in 7m53s
CI / Build (macos-latest) (push) Failing after 30s
CI / Build (ubuntu-latest) (push) Failing after 48s
CI / Distributed Training Tests (push) Has been skipped
CI / Clippy Check (push) Failing after 52s
CI / Format Check (push) Failing after 11s
CI / Test (macos-latest) (push) Has been skipped
CI / Test (ubuntu-latest) (push) Has been skipped
CI / Python Bindings (maturin) (macos-latest) (push) Has been skipped
CI / Python Bindings (maturin) (ubuntu-latest) (push) Has been skipped
CI / WASM Build + Size Check (push) Has been skipped
Documentation / Build User Guide (push) Successful in 9s
CI / Build CPU-Only (Explicit) (push) Failing after 1m10s
CI / CI Success (push) Failing after 0s
Documentation / Build API Documentation (push) Failing after 40s
Performance Benchmarks / Run Benchmarks (push) Successful in 7m53s
CI / Build (macos-latest) (push) Failing after 30s
CI / Build (ubuntu-latest) (push) Failing after 48s
CI / Distributed Training Tests (push) Has been skipped
CI / Clippy Check (push) Failing after 52s
Gap 2 — rtx-vision ViT bridge (jepa_vision_bridge.rs, 8 tests): - ViT::forward_features(): patch reps without classification head - ViT::encode_patch_indices(): shape-correct placeholder for GPU dispatch - RtxVisionJepaEncoder implementing JepaEncoder (vision-bridge feature) - From<&ViTConfig> for JepaViTConfig config conversion - rtx-vision added as optional dep; vision-bridge feature gate Gap 3 — WebDataset tar-shard reading (jepa_data.rs, +12 tests, 47 total): - parse_tar_bytes(): pure stdlib tar parser (512-byte block format) - read_webdataset_shard(): file reader with ShardLoadStats timing - WebDatasetRecord: key, image_bytes, label, extension - ShuffleBuffer: fixed-capacity reservoir sampling via LCG PRNG - JepaDataPipeline::from_filesystem(): validates paths, loads shards, builds pipeline Gap 5 — Training loop runner (jepa_runner.rs + examples/jepa_train.rs, 15 tests): - JepaRunConfig with TOML-style key=value parser - run_jepa_training(): full training loop (JepaTrainerV2, cosine LR, checkpointing) - JepaCheckpoint::save() writes JSON summary; load() stub - examples/jepa_train.rs: --config/--size/--steps/--dry-run CLI flags Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
e8a2036db4
commit
f487196367
@@ -112,6 +112,42 @@ impl ViT {
|
||||
.map_err(VisionError::from)
|
||||
}
|
||||
|
||||
/// Returns all patch representations `[seq_len+1, embed_dim]` (CLS token first,
|
||||
/// then patches). Does NOT apply the classification head. Suitable for feature
|
||||
/// extraction (e.g. JEPA encoders).
|
||||
pub fn forward_features(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let x = self.patch_embed.forward(x)?;
|
||||
let x = self.patch_embed.add_class_token(&x)?;
|
||||
let mut x = x;
|
||||
for block in &self.blocks {
|
||||
x = block.forward(&x)?;
|
||||
}
|
||||
self.norm.forward(&x)
|
||||
}
|
||||
|
||||
/// Encode specific patch indices (0-indexed, skipping the CLS token at
|
||||
/// position 0 in the feature sequence).
|
||||
///
|
||||
/// Returns a flattened `[n_patches, embed_dim]` f32 vector.
|
||||
///
|
||||
/// # Placeholder note
|
||||
///
|
||||
/// The returned values are currently **zeros of the correct shape**. Extracting
|
||||
/// arbitrary rows from a `Tensor` requires `narrow` + `to_vec()`, which depends
|
||||
/// on backend dispatch not yet wired for all targets. The shape contract
|
||||
/// (`n_patches * embed_dim` elements) is already correct and will be filled with
|
||||
/// real values once GPU tensor row-extraction is available (Batch 27).
|
||||
pub fn encode_patch_indices(&self, x: &Tensor, patch_indices: &[usize]) -> Result<Vec<f32>> {
|
||||
// Run full forward pass to get [n+1, embed_dim] features.
|
||||
let _features = self.forward_features(x)?;
|
||||
// patch i lives at position i+1 (position 0 is the CLS token).
|
||||
// Placeholder: return zeros with the correct shape until Tensor row
|
||||
// extraction (narrow + to_vec) is wired to the GPU backend.
|
||||
let n_patches = patch_indices.len();
|
||||
let d = self.config.embed_dim;
|
||||
Ok(vec![0.0f32; n_patches * d])
|
||||
}
|
||||
|
||||
fn extract_class_token(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let shape = x.shape().dims();
|
||||
match shape.len() {
|
||||
|
||||
Reference in New Issue
Block a user