Найти в Дзене
Один Rust не п...Rust

Использование мутабельного слайса вместо вектора в Rust для ML

Любое программное решение ограничено количеством памяти и CPU. Вектор позволяет хранить неограниченное количество информации, создавая тем самым угрозу работоспособности ПО. При этом рассчитать объем хранимой информации в векторе невозможно, по причине: Мутабельный слайс (&mut [T]) — это всего лишь "вид" на уже существующие данные. Он не выделяет память самостоятельно, а
только заимствует её у другой структуры (вектора, массива и т.д.). Сам объект слайса занимает на стеке фиксированный объём — два usize (указатель на данные и длина): use std::mem; fn main() { let vec = vec![1i32, 2, 3, 4, 5]; let slice: &mut [i32] = &mut vec.clone()[..]; // клонируем для мутабельности // Память, занимаемая самой структурой слайса на стеке let slice_struct_size = mem::size_of_val(&slice); println!("Размер структуры слайса на стеке: {} байт", slice_struct_size); // На 64-битной системе: 16 байт (8 байт указатель + 8 байт длина) } Всегда можно узнать, сколько памяти занимают данные, на которые ссылаетс
Оглавление

Почему нужно избегать использования векторов при работе с ML

Любое программное решение ограничено количеством памяти и CPU. Вектор позволяет хранить неограниченное количество информации, создавая тем самым угрозу работоспособности ПО. При этом рассчитать объем хранимой информации в векторе невозможно, по причине:

  • Внутри кода Rust можно посчитать только «полезную» нагрузку вектора, без служебных данных аллокатора.
  • Системные команды показывают всю память процесса, включая код, стек, кучу и другие библиотеки.
  • Аллокатор памяти в Rust может не всегда сразу возвращать освобождённую память обратно операционной системе, из-за чего RSS процесса может быть больше ожидаемого.

Почему нужно использовать мутабельный слайс при работе с ML

Мутабельный слайс (&mut [T]) — это всего лишь "вид" на уже существующие данные. Он не выделяет память самостоятельно, а
только заимствует её у другой структуры (вектора, массива и т.д.).

Сам объект слайса занимает на стеке фиксированный объём — два usize (указатель на данные и длина):

use std::mem;

fn main() {

let vec = vec![1i32, 2, 3, 4, 5];

let slice: &mut [i32] = &mut vec.clone()[..]; // клонируем для мутабельности

// Память, занимаемая самой структурой слайса на стеке

let slice_struct_size = mem::size_of_val(&slice);

println!("Размер структуры слайса на стеке: {} байт", slice_struct_size);

// На 64-битной системе: 16 байт (8 байт указатель + 8 байт длина)

}

Всегда можно узнать, сколько памяти занимают данные, на которые ссылается слайс:

fn calculate_slice_memory<T>(slice: &[T]) -> usize {

if slice.is_empty() {

return 0;

}

slice.len() * mem::size_of::<T>()

}

fn main() {

let mut vec = vec![1i64, 2, 3, 4, 5]; // i64 занимает 8 байт

let slice: &mut [i64] = &mut vec[1..4]; // слайс на элементы 2, 3, 4

let data_memory = calculate_slice_memory(slice);

println!("Память данных слайса: {} байт", data_memory);

// Вывод: 24 байта (3 элемента × 8 байт)

// Альтернативно:

let alt_calc = slice.len() * mem::size_of::<i64>();

println!("Альтернативный расчет: {} байт", alt_calc);

}

Слайс только заимствует память, работает с фактической длиной (len)

и позволяет расчитать:

  1. Структуру слайса: 16 байт на стеке (на 64-битных системах)
  2. Данные слайса: slice.len() * mem::size_of::<T>() байт в куче/стеке

Реализация нейронной сети с оптимизатором ADAM, где мутабельные слайсы позволяют изменять данные без копирования

use std::f64;

use rand; // Requires rand crate for dropout

fn main() {

// Toy dataset: XOR problem (batch size 4, features 2)

let inputs_flat = vec![0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0]; // Shape: (4, 2)

let targets_flat = vec![0.0, 1.0, 1.0, 0.0]; // Shape: (4, 1)

// Network architecture: input 2, hidden 4, hidden 4, output 1

let layer_sizes = vec![2, 4, 4, 1];

// Initialize weights and biases as flat Vec<f64>

let mut weights = vec![];

let mut biases = vec![];

let mut weight_shapes = vec![];

let mut bias_shapes = vec![];

let mut total_weights_len = 0;

let mut total_biases_len = 0;

for i in 1..layer_sizes.len() {

let prev = layer_sizes[i - 1];

let curr = layer_sizes[i];

let w_len = prev * curr;

let b_len = curr;

let mut w = vec![0.01; w_len]; // Simple init

let mut b = vec![0.0; b_len];

weights.extend_from_slice(&w);

biases.extend_from_slice(&b);

weight_shapes.push((prev, curr, total_weights_len));

bias_shapes.push((1, curr, total_biases_len));

total_weights_len += w_len;

total_biases_len += b_len;

}

// ADAM moments (first and second)

let mut weight_m = vec![0.0; total_weights_len];

let mut weight_v = vec![0.0; total_weights_len];

let mut bias_m = vec![0.0; total_biases_len];

let mut bias_v = vec![0.0; total_biases_len];

// Training params

let learning_rate = 0.01;

let beta1 = 0.9;

let beta2 = 0.999;

let epsilon = 1e-8;

let dropout_prob = 0.2;

let epochs = 1000;

let batch_size = 4; // Full batch

// Allocate activations, zs (post-activation), norms (hat_x), deltas (d_linear)

let mut activations: Vec<Vec<f64>> = layer_sizes.iter().map(|&size| vec![0.0; batch_size * size]).collect();

let mut zs: Vec<Vec<f64>> = layer_sizes[1..].iter().map(|&size| vec![0.0; batch_size * size]).collect();

let mut norms: Vec<Vec<f64>> = layer_sizes[1..].iter().map(|&size| vec![0.0; batch_size * size]).collect();

let mut deltas: Vec<Vec<f64>> = layer_sizes[1..].iter().map(|&size| vec![0.0; batch_size * size]).collect();

let mut means: Vec<Vec<f64>> = layer_sizes[1..].iter().map(|&size| vec![0.0; size]).collect(); // For batch norm

let mut vars: Vec<Vec<f64>> = layer_sizes[1..].iter().map(|&size| vec![1.0; size]).collect();

let mut t = 0; // ADAM timestep

for epoch in 0..epochs {

// Forward pass

forward(

&inputs_flat,

batch_size,

layer_sizes[0],

&weights,

&weight_shapes,

&biases,

&bias_shapes,

&mut activations,

&mut zs,

&mut norms,

&mut means,

&mut vars,

dropout_prob,

);

// Compute loss (cross-entropy)

let output_act = &activations[activations.len() - 1];

let mut loss = 0.0;

for i in 0..batch_size {

let y = targets_flat[i];

let p = output_act[i];

loss -= y * p.ln() + (1.0 - y) * (1.0 - p).ln();

}

loss /= batch_size as f64;

if epoch % 100 == 0 {

println!("Epoch {}: Loss = {:.4}", epoch, loss);

}

// Backprop

backprop(

&targets_flat,

batch_size,

&layer_sizes,

&activations,

&zs,

&norms,

&weights,

&weight_shapes,

&mut deltas,

&means,

&vars,

);

t += 1;

// Update weights with ADAM

update_weights(

learning_rate,

beta1,

beta2,

epsilon,

t,

&deltas,

&activations,

&layer_sizes,

&mut weights,

&weight_shapes,

&mut weight_m,

&mut weight_v,

&mut biases,

&bias_shapes,

&mut bias_m,

&mut bias_v,

);

}

// Inference on same data

forward(

&inputs_flat,

batch_size,

layer_sizes[0],

&weights,

&weight_shapes,

&biases,

&bias_shapes,

&mut activations,

&mut zs,

&mut norms,

&mut means,

&mut vars,

0.0, // No dropout

);

let predictions = &activations[activations.len() - 1];

println!("Predictions: {:?}", predictions);

}

// Forward propagation with batch norm and dropout

fn forward(

inputs: &[f64],

batch_size: usize,

input_size: usize,

weights: &[f64],

weight_shapes: &[(usize, usize, usize)],

biases: &[f64],

bias_shapes: &[(usize, usize, usize)],

activations: &mut [Vec<f64>],

zs: &mut [Vec<f64>],

norms: &mut [Vec<f64>],

means: &mut [Vec<f64>],

vars: &mut [Vec<f64>],

dropout_prob: f64,

) {

// Set input activation

activations[0].copy_from_slice(inputs);

for l in 0..weight_shapes.len() {

let (prev_size, curr_size, w_offset) = weight_shapes[l];

let (_, _, b_offset) = bias_shapes[l];

let w_slice = &weights[w_offset..w_offset + prev_size * curr_size];

let b_slice = &biases[b_offset..b_offset + curr_size];

let a_prev = &activations[l];

let mut z = &mut zs[l];

mat_mul(a_prev, batch_size, prev_size, w_slice, prev_size, curr_size, z);

add_bias(z, batch_size, curr_size, b_slice);

// Batch norm

batch_norm(z, batch_size, curr_size, &mut means[l], &mut vars[l], &mut norms[l]);

// Activation

if l < weight_shapes.len() - 1 {

relu(&norms[l], z);

} else {

sigmoid(&norms[l], z);

}

// Dropout (mutate in place)

if dropout_prob > 0.0 {

dropout(z, dropout_prob, z);

}

activations[l + 1].copy_from_slice(z);

}

}

// Backpropagation with full batch norm gradients

fn backprop(

targets: &[f64],

batch_size: usize,

layer_sizes: &[usize],

activations: &[Vec<f64>],

zs: &[Vec<f64>],

norms: &[Vec<f64>],

weights: &[f64],

weight_shapes: &[(usize, usize, usize)],

deltas: &mut [Vec<f64>],

means: &[Vec<f64>],

vars: &[Vec<f64>],

) {

let num_layers = layer_sizes.len();

for l in (0..num_layers - 1).rev() {

let curr_size = layer_sizes[l + 1];

let mut d_hat = vec![0.0; batch_size * curr_size];

if l == num_layers - 2 {

// Output layer: special for cross-entropy + sigmoid

let a_out = &activations[num_layers - 1];

for i in 0..batch_size * layer_sizes[num_layers - 1] {

d_hat[i] = a_out[i] - targets[i];

}

} else {

// Hidden layer: propagate from next

let next_size = layer_sizes[l + 2];

let (prev_size, _, w_offset) = weight_shapes[l + 1]; // w for next layer

let w_slice = &weights[w_offset..w_offset + prev_size * next_size];

let delta_next = &deltas[l];

let mut d_a = vec![0.0; batch_size * curr_size];

mat_mul_b_trans(delta_next, batch_size, next_size, w_slice, curr_size, next_size, &mut d_a);

let mut d_act = vec![0.0; batch_size * curr_size];

relu_deriv(&zs[l], &mut d_act);

elem_mul(&d_a, &d_act, &mut d_hat);

}

// Batch norm back for both

batch_norm_back(&d_hat, batch_size, curr_size, &means[l], &vars[l], &norms[l], &mut deltas[l]);

}

}

// Full batch norm backward pass

fn batch_norm_back(

dhat: &[f64],

batch_size: usize,

channels: usize,

mean: &[f64],

var: &[f64], // var here is sigma = sqrt(variance + eps)

hat: &[f64], // norms, hat_x

out: &mut [f64],

) {

assert_eq!(dhat.len(), batch_size * channels);

assert_eq!(hat.len(), batch_size * channels);

assert_eq!(out.len(), batch_size * channels);

assert_eq!(mean.len(), channels);

assert_eq!(var.len(), channels);

for c in 0..channels {

let sigma = var[c];

let mut sum_dhat = 0.0;

let mut sum_dhat_hat = 0.0;

for b in 0..batch_size {

let idx = b * channels + c;

sum_dhat += dhat[idx];

sum_dhat_hat += dhat[idx] * hat[idx];

}

let dl_dvar = -0.5 * sum_dhat_hat / (sigma * sigma);

let dl_dmean = -sum_dhat / sigma;

let n = batch_size as f64;

for b in 0..batch_size {

let idx = b * channels + c;

let h = hat[idx];

out[idx] = dhat[idx] / sigma + dl_dmean / n + dl_dvar * 2.0 * h * sigma / n;

}

}

}

// Update weights and biases with ADAM

fn update_weights(

alpha: f64,

beta1: f64,

beta2: f64,

epsilon: f64,

t: usize,

deltas: &[Vec<f64>],

activations: &[Vec<f64>],

layer_sizes: &[usize],

weights: &mut [f64],

weight_shapes: &[(usize, usize, usize)],

weight_m: &mut [f64],

weight_v: &mut [f64],

biases: &mut [f64],

bias_shapes: &[(usize, usize, usize)],

bias_m: &mut [f64],

bias_v: &mut [f64],

) {

let batch_size = activations[0].len() / layer_sizes[0];

for l in 0..weight_shapes.len() {

let (prev_size, curr_size, w_offset) = weight_shapes[l];

let (_, _, b_offset) = bias_shapes[l];

let delta = &deltas[l]; // d_linear

let a_prev = &activations[l];

let mut dw = vec![0.0; prev_size * curr_size];

mat_mul_transpose(a_prev, batch_size, prev_size, delta, batch_size, curr_size, &mut dw);

// Average gradient over batch

for val in dw.iter_mut() {

*val /= batch_size as f64;

}

let mut w_slice = &mut weights[w_offset..w_offset + prev_size * curr_size];

let mut wm_slice = &mut weight_m[w_offset..w_offset + prev_size * curr_size];

let mut wv_slice = &mut weight_v[w_offset..w_offset + prev_size * curr_size];

update_with_adam(&dw, alpha, beta1, beta2, epsilon, t, wm_slice, wv_slice, w_slice);

// Bias update (mean over batch)

let mut db = vec![0.0; curr_size];

for i in 0..curr_size {

for b in 0..batch_size {

db[i] += delta[b * curr_size + i];

}

db[i] /= batch_size as f64;

}

let mut b_slice = &mut biases[b_offset..b_offset + curr_size];

let mut bm_slice = &mut bias_m[b_offset..b_offset + curr_size];

let mut bv_slice = &mut bias_v[b_offset..b_offset + curr_size];

update_with_adam(&db, alpha, beta1, beta2, epsilon, t, bm_slice, bv_slice, b_slice);

}

}

fn update_with_adam(grad: &[f64], alpha: f64, beta1: f64, beta2: f64, epsilon: f64, t: usize, m: &mut [f64], v: &mut [f64], param: &mut [f64]) {

for (i, &g) in grad.iter().enumerate() {

m[i] = beta1 * m[i] + (1.0 - beta1) * g;

v[i] = beta2 * v[i] + (1.0 - beta2) * g * g;

let m_hat = m[i] / (1.0 - beta1.powi(t as i32));

let v_hat = v[i] / (1.0 - beta2.powi(t as i32));

param[i] -= alpha * m_hat / (v_hat.sqrt() + epsilon);

}

}

// Matrix multiplication A * B (row-major)

fn mat_mul(a: &[f64], a_rows: usize, a_cols: usize, b: &[f64], b_rows: usize, b_cols: usize, out: &mut [f64]) {

assert_eq!(a_cols, b_rows);

assert_eq!(a.len(), a_rows * a_cols);

assert_eq!(b.len(), b_rows * b_cols);

assert_eq!(out.len(), a_rows * b_cols);

out.fill(0.0);

for i in 0..a_rows {

for j in 0..b_cols {

let mut sum = 0.0;

for k in 0..a_cols {

sum += a[i * a_cols + k] * b[k * b_cols + j];

}

out[i * b_cols + j] = sum;

}

}

}

// Matrix multiplication A^T * B (for gradients, dw = a_prev^T * delta)

fn mat_mul_transpose(a: &[f64], a_rows: usize, a_cols: usize, b: &[f64], b_rows: usize, b_cols: usize, out: &mut [f64]) {

// A^T * B, A^T is a_cols x a_rows, B a_rows x b_cols, out a_cols x b_cols

assert_eq!(a_rows, b_rows);

assert_eq!(a.len(), a_rows * a_cols);

assert_eq!(b.len(), b_rows * b_cols);

assert_eq!(out.len(), a_cols * b_cols);

out.fill(0.0);

for i in 0..a_cols {

for j in 0..b_cols {

let mut sum = 0.0;

for k in 0..a_rows {

sum += a[k * a_cols + i] * b[k * b_cols + j];

}

out[i * b_cols + j] = sum;

}

}

}

// Matrix multiplication A * B^T (for backprop, delta * w^T)

fn mat_mul_b_trans(a: &[f64], a_rows: usize, a_cols: usize, b: &[f64], b_rows: usize, b_cols: usize, out: &mut [f64]) {

// A * B^T, B^T is b_cols x b_rows, out a_rows x b_rows

assert_eq!(a_cols, b_cols);

assert_eq!(a.len(), a_rows * a_cols);

assert_eq!(b.len(), b_rows * b_cols);

assert_eq!(out.len(), a_rows * b_rows);

out.fill(0.0);

for i in 0..a_rows {

for j in 0..b_rows {

let mut sum = 0.0;

for k in 0..a_cols {

sum += a[i * a_cols + k] * b[j * b_cols + k];

}

out[i * b_rows + j] = sum;

}

}

}

// Add bias (broadcast)

fn add_bias(z: &mut [f64], rows: usize, cols: usize, bias: &[f64]) {

for row in 0..rows {

for j in 0..cols {

z[row * cols + j] += bias[j];

}

}

}

// Element-wise multiply

fn elem_mul(a: &[f64], b: &[f64], out: &mut [f64]) {

assert_eq!(a.len(), b.len());

assert_eq!(a.len(), out.len());

for i in 0..a.len() {

out[i] = a[i] * b[i];

}

}

// Sigmoid (from input to out)

fn sigmoid(input: &[f64], out: &mut [f64]) {

for (i, &val) in input.iter().enumerate() {

out[i] = 1.0 / (1.0 + (-val).exp());

}

}

// ReLU (from input to out)

fn relu(input: &[f64], out: &mut [f64]) {

for (i, &val) in input.iter().enumerate() {

out[i] = val.max(0.0);

}

}

// ReLU derivative (can use on post-relu since same)

fn relu_deriv(z: &[f64], out: &mut [f64]) {

for (i, &val) in z.iter().enumerate() {

out[i] = if val > 0.0 { 1.0 } else { 0.0 };

}

}

// Dropout in place (mask with prob)

fn dropout(act: &mut [f64], prob: f64, out: &mut [f64]) {

for (i, val) in act.iter_mut().enumerate() {

if rand::random::<f64>() < prob {

out[i] = 0.0;

} else {

out[i] = *val / (1.0 - prob);

}

}

}

// Batch norm forward (compute mean/var, normalize to norm)

fn batch_norm(

z: &[f64],

batch_size: usize,

channels: usize,

mean: &mut [f64],

var: &mut [f64],

norm: &mut [f64],

) {

// Compute mean

for c in 0..channels {

mean[c] = 0.0;

for b in 0..batch_size {

mean[c] += z[b * channels + c];

}

mean[c] /= batch_size as f64;

}

// Compute variance

for c in 0..channels {

var[c] = 0.0;

for b in 0..batch_size {

let diff = z[b * channels + c] - mean[c];

var[c] += diff * diff;

}

var[c] /= batch_size as f64;

var[c] = (var[c] + 1e-5).sqrt();

}

// Normalize

for b in 0..batch_size {

for c in 0..channels {

norm[b * channels + c] = (z[b * channels + c] - mean[c]) / var[c];

}

}

}

  • Основная функция организует обучение. Изменяемые срезы используются в вызовах функций более высокого уровня и локальных вычислениях.
  • Переменные, let mut weights = vec![];", "let mut weight_m = vec![0.0; total_weights_len]; создают изменяемые векторы, которые затем нарезаются изменяемым образом.
  • В forward(&inputs_flat, ..., &mut activations, &mut zs, &mut norms, &mut mean, &mut vars, ...); изменяемые срезы Vec<f64> (&mut [Vec<f64>]) передаются для послойных данных. let mut loss = 0.0; это простой изменяемый скаляр, но он связан с неизменяемым доступом к срезу через &activations[...] для чтения выходных данных.
    backprop(&targets_flat, ..., &mut deltas, ...); использует &mut [Vec<f64>] для обновления дельт ошибок между слоями.
  • update_weights(learning_rate, ..., &mut weights, ..., &mut weight_m, &mut weight_v, &mut biases, ..., &mut bias_m, &mut bias_v); передает &mut [f64] для параметров и моментов ADAM. Это обеспечивает прямую мутацию обучаемых параметров сети. Изменяемость здесь гарантирует, что итерации обучения могут обновлять состояние без перераспределения больших векторов в каждой эпохе.
  • Такие параметры, как activations: &mut [Vec<f64>], zs: &mut [Vec<f64>], norms: &mut [Vec<f64>], meanings: &mut [Vec<f64>], vars: &mut [Vec<f64>] — это изменяемые срезы векторов, позволяющие выполнять мутации, специфичные для слоя.
  • let mut z = &mut zs[l]; создаёт изменяемую ссылку на отдельный Vec<f64>, который рассматривается как срез для дальнейших операций.
    mat_mul(..., z); и add_bias(z, ...); передача z: &mut [f64] (неявно через deref) для мутации выходных данных.
  • batch_norm(z, ..., &mut mean[l], &mut vars[l], &mut norms[l]); использует изменяемые подсрезы для обновления статистики.
  • relu(&norms[l], z); и выпадение (z, ..., z); мутация z на месте. Это обеспечивает эффективные прямые проходы за счет изменения предварительно выделенных буферов.
  • deltas: &mut [Vec<f64>] для обновления ошибок.
  • mat_mul_b_trans(..., &mut d_a);, relu_deriv(..., &mut d_act);, elem_mul(..., &mut d_hat);, batch_norm_back(..., &mut deltas[l]); используйте &mut [f64] для временных векторов, что позволяет распространять градиенты на месте. Мутабильность поддерживает необходимость обратного прохода накапливать и изменять дельты слой за слоем.
  • out: &mut [f64] для вывода нормализованных градиентов.
  • Множественные &mut [f64] для весов, смещений и моментов.
    mat_mul_transpose(..., &mut dw);, let mut w_slice = &mut weights[...]; и т.д., создают и изменяют подсрезы для обновлений, специфичных для параметров.
  • m: &mut [f64], v: &mut [f64], param: &mut [f64] для обновлений моментов и параметров.
  • выход: &mut [f64] для накопления результатов.
  • выход: &mut [f64]. Эти функции мутируют выходные данные во вложенных циклах, что является основой линейной алгебры в нейронных сетях.
  • z: &mut [f64] для широковещательных дополнений. Мутирует предактивационный срез.
  • выход: &mut [f64] для поэлементных произведений. Используется в обратном распространении ошибки для цепного правила.
  • &mut [f64] для вычисления активаций/производных на месте.
  • act: &mut [f64], out: &mut [f64] (часто один и тот же срез). Мутирует активации для регуляризации.
  • mean: &mut [f64], var: &mut [f64], norm: &mut [f64]. Вычисляет и мутирует статистику для нормализации.