ferrite-dl 0.2.0

Deep learning library written in pure Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use crate::tensor::*;
use std::sync::{Arc, RwLock};
use rand::seq::SliceRandom;
use rand::thread_rng;

// Enum to specify which dataset split to access
pub enum Split {
  Train,
  Validation,
  Test,
}

pub struct DataLoader {
  x_data: Arc<RwLock<Tensor>>,
  y_data: Arc<RwLock<Tensor>>,
  batch_size: usize,
  train_split: f32,
  test_split: f32,
  validation_split: f32,
  train_size: usize,
  test_size: usize,
  validation_size: usize,
}

impl DataLoader {
  /// Creates a new DataLoader instance, computing split sizes based on ratios.
  pub fn new(
    x_data: Arc<RwLock<Tensor>>,
    y_data: Arc<RwLock<Tensor>>,
    batch_size: usize,
    train_split: f32,
    test_split: f32,
    validation_split: f32,
  ) -> Self {
    // Ensure splits sum to 1.0 (with small tolerance for floating-point errors)
    let total_split = train_split + test_split + validation_split;
    assert!(
      (total_split - 1.0).abs() < 1e-6,
      "Split ratios must sum to 1.0"
    );

    // Get the total number of samples from x_data
    let x_data_read = x_data.read().unwrap();
    let n_samples = x_data_read.shape()[0];

    // Compute sizes for each split
    let train_size = (n_samples as f32 * train_split).floor() as usize;
    let test_size = (n_samples as f32 * test_split).floor() as usize;
    let validation_size = n_samples - train_size - test_size;

    Self {
      x_data,
      y_data,
      batch_size,
      train_split,
      test_split,
      validation_split,
      train_size,
      test_size,
      validation_size,
    }
  }

  /// Returns an iterator over batches for the specified split.
  /// If `shuffle` is true, the indices are shuffled before batching.
  pub fn batches(&self, split: Split, shuffle: bool) -> BatchIterator {
    // Determine the index range based on the split
    let (start, end) = match split {
      Split::Train => (0, self.train_size),
      Split::Validation => (self.train_size, self.train_size + self.validation_size),
      Split::Test => (
        self.train_size + self.validation_size,
        self.train_size + self.validation_size + self.test_size,
      ),
    };

    // Collect indices for the split
    let mut indices: Vec<usize> = (start..end).collect();

    // Shuffle indices if requested
    if shuffle {
      let mut rng = thread_rng();
      indices.shuffle(&mut rng);
    }

    BatchIterator {
      x_data: self.x_data.clone(),
      y_data: self.y_data.clone(),
      indices,
      batch_size: self.batch_size,
      current: 0,
    }
  }
}

/// Iterator over batches of data from the DataLoader.
pub struct BatchIterator {
    x_data: Arc<RwLock<Tensor>>,
    y_data: Arc<RwLock<Tensor>>,
    indices: Vec<usize>,
    batch_size: usize,
    current: usize,
}

impl Iterator for BatchIterator {
  type Item = (Tensor, Tensor);

  fn next(&mut self) -> Option<Self::Item> {
    if self.current >= self.indices.len() {
      None
    } else {
      // Calculate batch boundaries
      let start = self.current;
      let end = (start + self.batch_size).min(self.indices.len());
      let batch_indices = &self.indices[start..end];
      self.current = end;

      // Access data with read locks
      let x_data_read = self.x_data.read().unwrap();
      let y_data_read = self.y_data.read().unwrap();

      // Select batch data
      let x_batch = x_data_read.index_select(0, batch_indices);
      let y_batch = y_data_read.index_select(0, batch_indices);

      Some((x_batch, y_batch))
    }
  }
}