Class: GRX::Data::DataLoader

Inherits:
Object
  • Object
show all
Includes:
Enumerable
Defined in:
lib/grx/data.rb

Overview

================================================================

DataLoader — Mini-batch iterator with optional shuffling

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(dataset, batch_size: 32, shuffle: true) ⇒ DataLoader

Returns a new instance of DataLoader.



53
54
55
56
57
 
# File 'lib/grx/data.rb', line 53

def initialize(dataset, batch_size: 32, shuffle: true)
  @dataset    = dataset
  @batch_size = batch_size
  @shuffle    = shuffle
end

Instance Attribute Details

#batch_size ⇒ Object (readonly)

Returns the value of attribute batch_size.



51
52
53
 
# File 'lib/grx/data.rb', line 51

def batch_size
  @batch_size
end

#dataset ⇒ Object (readonly)

Returns the value of attribute dataset.



51
52
53
 
# File 'lib/grx/data.rb', line 51

def dataset
  @dataset
end

#shuffle ⇒ Object (readonly)

Returns the value of attribute shuffle.



51
52
53
 
# File 'lib/grx/data.rb', line 51

def shuffle
  @shuffle
end

Instance Method Details

#each ⇒ Object 



59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
 
# File 'lib/grx/data.rb', line 59

def each
  return to_enum(:each) unless block_given?

  indices = (0...@dataset.size).to_a
  indices.shuffle! if @shuffle

  indices.each_slice(@batch_size) do |batch_indices|
    batch_samples = batch_indices.map { |i| @dataset[i] }
    num_tensors = batch_samples.first.size

    batched = (0...num_tensors).map do |t_idx|
      slices = batch_samples.map { |sample| sample[t_idx].to_a }
      flat_data = slices.flatten
      sample_shape = batch_samples.first[t_idx].shape
      batch_dim = batch_indices.size
      target_shape = [batch_dim] + sample_shape[1..]
      Tensor.create(flat_data, target_shape)
    end

    yield(*batched)
  end
end

#size ⇒ Object 



82
83
84
 
# File 'lib/grx/data.rb', line 82

def size
  (@dataset.size.to_f / @batch_size).ceil
end