use crate::tensor::{Tensor, shape};
pub trait Dataset: Send + Sync {
fn len(&self) -> usize;
fn input_shape(&self) -> Vec<usize>;
fn target_shape(&self) -> Vec<usize>;
fn write(&self, index: usize, input: &mut [f32], target: &mut [f32]);
fn is_empty(&self) -> bool {
self.len() == 0
}
fn input_size(&self) -> usize {
shape::numel(&self.input_shape())
}
fn target_size(&self) -> usize {
shape::numel(&self.target_shape())
}
}
pub struct TensorDataset {
inputs: Vec<f32>,
targets: Vec<f32>,
input_shape: Vec<usize>,
target_shape: Vec<usize>,
len: usize,
}
impl TensorDataset {
pub fn new(inputs: &Tensor, targets: &Tensor) -> TensorDataset {
assert_eq!(
inputs.dim(0),
targets.dim(0),
"TensorDataset: {} inputs but {} targets",
inputs.dim(0),
targets.dim(0)
);
TensorDataset {
len: inputs.dim(0),
input_shape: item_shape(inputs),
target_shape: item_shape(targets),
inputs: inputs.to_vec(),
targets: targets.to_vec(),
}
}
}
impl Dataset for TensorDataset {
fn len(&self) -> usize {
self.len
}
fn input_shape(&self) -> Vec<usize> {
self.input_shape.clone()
}
fn target_shape(&self) -> Vec<usize> {
self.target_shape.clone()
}
fn write(&self, index: usize, input: &mut [f32], target: &mut [f32]) {
copy_item(&self.inputs, index, input);
copy_item(&self.targets, index, target);
}
}
fn item_shape(t: &Tensor) -> Vec<usize> {
if t.ndim() <= 1 {
vec![1]
} else {
t.shape()[1..].to_vec()
}
}
fn copy_item(source: &[f32], index: usize, out: &mut [f32]) {
let start = index * out.len();
out.copy_from_slice(&source[start..start + out.len()]);
}