Implementation:LaurentMazare Tch rs Dataset Train Iter

Knowledge Sources	tch-rs
Domains	Deep_Learning, Data_Loading
Last Updated	2026-02-08 14:00 GMT

Overview

Concrete tool for creating mini-batch iterators over vision datasets provided by the tch vision module.

Description

Dataset::train_iter creates an Iter2 iterator that yields (images, labels) tensor pairs of the specified batch size from the training split. The Iter2 type supports method chaining with .shuffle() for random ordering, .to_device() for GPU transfer, and .return_smaller_last_batch() to include partial final batches.

Usage

Use in training loops to iterate over the training dataset. Chain with .shuffle() and .to_device() as needed.

Code Reference

Source Location

• Repository: tch-rs
• File: src/vision/dataset.rs
• Lines: 16-18

Signature

impl Dataset {
    pub fn train_iter(&self, batch_size: i64) -> Iter2
}

Import

use tch::vision::mnist;  // or imagenet
// Dataset is returned by load_dir / load_from_dir

I/O Contract

Inputs

Name	Type	Required	Description
batch_size	i64	Yes	Number of samples per mini-batch

Outputs

Name	Type	Description
Iter2	Iter2	Iterator yielding (Tensor, Tensor) pairs; supports .shuffle(), .to_device()

Usage Examples

let m = tch::vision::mnist::load_dir("data/mnist")?;

for (images, labels) in m.train_iter(64).shuffle().to_device(device) {
    let logits = net.forward(&images);
    let loss = logits.cross_entropy_for_logits(&labels);
    opt.backward_step(&loss);
}

Related Pages

Implements Principle

• Principle:LaurentMazare_Tch_rs_Dataset_Iteration