Struct DistributedSampler 

pub struct DistributedSampler { /* private fields */ }

Distributed sampler for balanced data distribution

Unlike DistributedWrapper which wraps an existing sampler, DistributedSampler is a standalone sampler designed specifically for distributed training. It provides more control over the distribution strategy.

Examples

use torsh_data::sampler::{DistributedSampler, Sampler};

// Process 1 of 4 processes, working with dataset of size 100
let sampler = DistributedSampler::new(100, 4, 1, true).with_generator(42);

let indices: Vec<usize> = sampler.iter().collect();
assert_eq!(indices.len(), 25); // 100 / 4 = 25 samples per process

Implementations

impl DistributedSampler

pub fn new( dataset_size: usize, num_replicas: usize, rank: usize, shuffle: bool, ) -> Self

Create a new distributed sampler

Arguments

• dataset_size - Total size of the dataset
• num_replicas - Total number of processes
• rank - Current process rank (0-based)
• shuffle - Whether to shuffle indices

Panics

Panics if dataset_size is 0, num_replicas is 0, or rank >= num_replicas

pub fn with_generator(self, seed: u64) -> Self

Set the random number generator seed

pub fn with_drop_last(self, drop_last: bool) -> Self

Set whether to drop the last samples to make dataset evenly divisible

pub fn dataset_size(&self) -> usize

Get the dataset size

pub fn num_replicas(&self) -> usize

Get the number of replicas

pub fn rank(&self) -> usize

Get the current rank

pub fn shuffle(&self) -> bool

Check if shuffling is enabled

pub fn drop_last(&self) -> bool

Check if dropping last samples

pub fn generator(&self) -> Option<u64>

Get the generator seed if set

Trait Implementations

impl Clone for DistributedSampler

fn clone(&self) -> DistributedSampler

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for DistributedSampler

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Sampler for DistributedSampler

type Iter = SamplerIterator

Iterator type returned by the sampler

fn iter(&self) -> Self::Iter

Create an iterator over indices

fn len(&self) -> usize

Total number of samples that will be yielded

fn is_empty(&self) -> bool

Check if sampler is empty

fn into_batch_sampler( self, batch_size: usize, drop_last: bool, ) -> BatchingSampler<Self>

where Self: Sized,

Convert this sampler into a batch sampler

fn into_distributed( self, num_replicas: usize, rank: usize, ) -> DistributedWrapper<Self>

where Self: Sized,

Create a distributed version of this sampler