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

Использование usize и isize в ML на Rust

При работе с машинным обучением на Rust типы usize и isize играют ключевую роль в обеспечении эффективности, безопасности памяти и переносимости кода. Эти типы особенно важны при обработке больших данных, индексации тензоров и управлении памятью. Вот как их можно эффективно использовать в контексте ML: В ML часто работают с многомерными массивами и тензорами. Тип usize
является стандартным для индексации, так как его размер соответствует
архитектуре платформы (32 или 64 бита), что гарантирует покрытие всего
адресного пространства. Например: let tensor = vec![1.0, 2.0, 3.0, 4.0]; let index: usize = 2; let value = tensor[index]; // Доступ к элементу тензора При определении архитектуры нейронных сетей (например, линейных слоёв) размерности входов и выходов задаются с помощью usize, так как они представляют неотрицательные значения (количество нейронов, каналов и т.д.): struct LinearLayer { weights: Tensor, // Размерности: [input_size, output_size] bias: Tensor, // Размерность: [o
Оглавление

При работе с машинным обучением на Rust типы usize и isize играют ключевую роль в обеспечении эффективности, безопасности памяти и переносимости кода. Эти типы особенно важны при обработке больших данных, индексации тензоров и управлении памятью. Вот как их можно эффективно использовать в контексте ML:

  • Обеспечить безопасность памяти за счет строгой типизации и проверок границ.
  • Повысить переносимость кода между архитектурами.
  • Эффективно работать с большими данными благодаря соответствию размеру указателя.
  • Интегрироваться с низкоуровневыми библиотеками (например, LibTorch через tch-rs).
  • Проверить наличие отрицательных индексов: В ML индексы обычно неотрицательные. Если нужны относительные смещения, используйте isize, но преобразуйте в usize перед индексацией.
  • Проверить переполнение: При работе с большими данными на 64-битных системах используйте проверенные операции (например, checked_add, saturating_sub).

📌 1. Индексация тензоров и коллекций

В ML часто работают с многомерными массивами и тензорами. Тип usize
является стандартным для индексации, так как его размер соответствует
архитектуре платформы (32 или 64 бита), что гарантирует покрытие всего
адресного пространства. Например:

let tensor = vec![1.0, 2.0, 3.0, 4.0];

let index: usize = 2;

let value = tensor[index]; // Доступ к элементу тензора

  • Почему usize? Индексы не могут быть отрицательными, поэтому usize идеален для этой цели.
  • Пример из ML: В крейте tch-rs (Rust-обёртка для LibTorch) индексы для доступа к элементам тензоров используют usize.

📌 2. Задание размерностей слоёв нейросетей

При определении архитектуры нейронных сетей (например, линейных слоёв) размерности входов и выходов задаются с помощью usize, так как они представляют неотрицательные значения (количество нейронов, каналов и т.д.):

struct LinearLayer {

weights: Tensor, // Размерности: [input_size, output_size]

bias: Tensor, // Размерность: [output_size]

}

impl LinearLayer {

fn new(input_size: usize, output_size: usize) -> Self {

// Инициализация весов и смещений

}

}

Пример из поиска: размерности для линейного и сверточных слоёх задаются как usize (например, in_channels: usize, out_channels: usize).

📌 3. Работа с памятью и буферами

Тип usize используется для управления памятью, например, для определения размеров буферов, выделения памяти и работы с указателями:

fn allocate_buffer(size: usize) -> Vec<f32> {

vec![0.0; size] // Выделение буфера заданного размера

}

Важно: В ML-фреймворках все тензоры хранятся в центральном хранилище (Memory), где индексы параметров имеют тип usize.

📌 4. Обработка смещений и шагов (strides)

Для операций, требующих отрицательных значений (например, смещения или относительные индексы), можно использовать isize. Однако при индексации его обычно преобразуют в usize после проверки:

fn compute_offset(offset: isize, base: usize) -> Option<usize> {

if offset >= 0 {

Some(base + offset as usize)

} else {

let abs_offset = (-offset) as usize;

base.checked_sub(abs_offset) // Проверка на переполнение

}

}

Пример: В сверточных слоях при вычислении паддинга или смещений может потребоваться isize, но для конечной индексации используется usize.

📌 5. Безопасное преобразование типов

В ML часто приходится работать с данными разной разрядности (например, u32, i64). Для преобразования в usize/isize используйте явное приведение с проверками:

let x: u32 = 1000;

let idx: usize = x as usize; // Явное преобразование

// Безопасное преобразование с проверкой границ

let y: i64 = -10;

if y >= 0 {

let idx: usize = y as usize;

} else {

eprintln!("Ошибка: отрицательный индекс!");

}

Важно: Преобразование из i64 в usize на 32-битной платформе может привести к усечению данных. Используйте методы типа try_into() для безопасного преобразования.

Пример использования в оптимизации вычислений

В низкоуровневых оптимизациях (например, в бинарном поиске для обработки данных) usize обеспечивает эффективную индексацию. Например, в оптимизированной версии binary_search_by из Rust 1.52+ используется usize для индексов:

let s = [0, 1, 1, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55];

let seek = 13;

assert_eq!(s.binary_search_by(|probe| probe.cmp(&seek)), Ok(9));

// Возвращает usize

Пример использования usize в контексте нейронных сетей:

  1. Индексация многомерных массивов: usize используется для доступа к весам и смещениям
  2. Проверка границ: Все операции с индексами включают проверку выхода за границы
  3. Управление памятью: Расчет использования памяти и кэширование
  4. SIMD оптимизация: Использование векторных операций для ускорения вычислений
  5. Обработка ошибок: Детальные сообщения об ошибках с использованием usize в диагностике
  6. Квантование весов: Специализированная реализация для f32
  7. Пакетная обработка: Поддержка обработки нескольких входных данных одновременно
  8. Разреженные вычисления: Оптимизация для работы только с активными нейронами

use std::ops::{Add, Mul};

use std::simd::{f32x4, SimdFloat};

use std::collections::HashMap;

use std::hash::Hash;

use thiserror::Error;

#[derive(Error, Debug, Clone)]

pub enum LinearLayerError {

#[error("Размерность входа {input_size} не совпадает с ожидаемой {expected}")]

InputSizeMismatch { input_size: usize, expected: usize },

#[error("Индекс {index} выходит за границы размерности {dimension}")]

IndexOutOfBounds { index: usize, dimension: usize },

#[error("Неверная инициализация: {reason}")]

InitializationError { reason: String },

}

#[derive(Debug, Clone)]

pub struct LinearLayer<T> {

weights: Vec<Vec<T>>,

biases: Vec<T>,

input_size: usize,

output_size: usize,

weight_cache: HashMap<(usize, usize), T>,

use_simd: bool,

}

impl<T> LinearLayer<T>

where

T: Copy + Default + Add<Output = T> + Mul<Output = T> + From<f32> + Into<f64>,

{

pub fn new(input_size: usize, output_size: usize) -> Result<Self, LinearLayerError> {

if input_size == 0 || output_size == 0 {

return Err(LinearLayerError::InitializationError {

reason: "Размерности должны быть больше нуля".to_string(),

});

}

// Инициализация весов методом Кайминга (He initialization)

let std_dev = (2.0 / input_size as f64).sqrt();

let mut weights = Vec::with_capacity(output_size);

for i in 0..output_size {

let mut neuron_weights = Vec::with_capacity(input_size);

for j in 0..input_size {

let weight = T::from((rand::random::<f64>() - 0.5) * std_dev * 2.0) as T;

neuron_weights.push(weight);

}

weights.push(neuron_weights);

}

// Инициализация смещений нулями

let biases = vec![T::default(); output_size];

Ok(Self {

weights,

biases,

input_size,

output_size,

weight_cache: HashMap::new(),

use_simd: false,

})

}

pub fn enable_simd(&mut self, enable: bool) {

self.use_simd = enable && self.input_size % 4 == 0;

}

pub fn forward(&mut self, input: &[T]) -> Result<Vec<T>, LinearLayerError> {

if input.len() != self.input_size {

return Err(LinearLayerError::InputSizeMismatch {

input_size: input.len(),

expected: self.input_size,

});

}

let mut output = Vec::with_capacity(self.output_size);

if self.use_simd && std::any::TypeId::of::<T>() == std::any::TypeId::of::<f32>() {

// Использование SIMD для оптимизации вычислений

unsafe { self.forward_simd(input, &mut output)? };

} else {

// Стандартная реализация

for i in 0..self.output_size {

let mut sum = T::default();

for j in 0..self.input_size {

sum = sum + input[j] * self.weights[i][j];

}

output.push(sum + self.biases[i]);

}

}

Ok(output)

}

#[target_feature(enable = "avx2")]

unsafe fn forward_simd(&self, input: &[T], output: &mut Vec<T>) -> Result<(), LinearLayerError> {

// Приведение типов для SIMD операций

if std::any::TypeId::of::<T>() != std::any::TypeId::of::<f32>() {

return Ok(());

}

let input_f32: &[f32] = std::mem::transmute(input);

let biases_f32: &[f32] = std::mem::transmute(&self.biases[..]);

for i in 0..self.output_size {

let weights_f32: &[f32] = std::mem::transmute(&self.weights[i][..]);

let mut sum = f32x4::splat(0.0);

// Обработка по 4 элемента за раз

for j in (0..self.input_size).step_by(4) {

let input_simd = f32x4::from_slice(&input_f32[j..]);

let weights_simd = f32x4::from_slice(&weights_f32[j..]);

sum += input_simd * weights_simd;

}

// Суммирование компонент SIMD вектора

let mut result = sum.reduce_sum();

result += biases_f32[i];

output.push(std::mem::transmute_copy(&result));

}

Ok(())

}

pub fn get_weight(&mut self, i: usize, j: usize) -> Result<T, LinearLayerError> {

if i >= self.output_size {

return Err(LinearLayerError::IndexOutOfBounds {

index: i,

dimension: self.output_size,

});

}

if j >= self.input_size {

return Err(LinearLayerError::IndexOutOfBounds {

index: j,

dimension: self.input_size,

});

}

// Кэширование для быстрого доступа к часто используемым весам

let key = (i, j);

if let Some(&weight) = self.weight_cache.get(&key) {

return Ok(weight);

}

let weight = self.weights[i][j];

self.weight_cache.insert(key, weight);

Ok(weight)

}

pub fn set_weight(&mut self, i: usize, j: usize, value: T) -> Result<(), LinearLayerError> {

if i >= self.output_size {

return Err(LinearLayerError::IndexOutOfBounds {

index: i,

dimension: self.output_size,

});

}

if j >= self.input_size {

return Err(LinearLayerError::IndexOutOfBounds {

index: j,

dimension: self.input_size,

});

}

self.weights[i][j] = value;

self.weight_cache.insert((i, j), value);

Ok(())

}

pub fn apply_to_indices<F>(&mut self, indices: &[(usize, usize)], mut func: F) -> Result<(), LinearLayerError>

where

F: FnMut(T) -> T,

{

for &(i, j) in indices {

if i >= self.output_size || j >= self.input_size {

continue;

}

let new_value = func(self.weights[i][j]);

self.weights[i][j] = new_value;

self.weight_cache.insert((i, j), new_value);

}

Ok(())

}

pub fn sparse_forward(&self, input: &[T], active_neurons: &[usize]) -> Result<Vec<T>, LinearLayerError> {

if input.len() != self.input_size {

return Err(LinearLayerError::InputSizeMismatch {

input_size: input.len(),

expected: self.input_size,

});

}

let mut output = Vec::with_capacity(active_neurons.len());

for &neuron_idx in active_neurons {

if neuron_idx >= self.output_size {

return Err(LinearLayerError::IndexOutOfBounds {

index: neuron_idx,

dimension: self.output_size,

});

}

let mut sum = T::default();

for j in 0..self.input_size {

sum = sum + input[j] * self.weights[neuron_idx][j];

}

output.push(sum + self.biases[neuron_idx]);

}

Ok(output)

}

pub fn memory_usage(&self) -> usize {

let weights_size = self.output_size * self.input_size * std::mem::size_of::<T>();

let biases_size = self.output_size * std::mem::size_of::<T>();

let cache_size = self.weight_cache.len() * (std::mem::size_of::<(usize, usize)>() + std::mem::size_of::<T>());

weights_size + biases_size + cache_size

}

}

// Реализация трейта для совместимости с автоматическим дифференцированием

impl<T> crate::autodiff::Differentiable for LinearLayer<T> {

fn parameters(&self) -> Vec<crate::autodiff::Parameter> {

let mut params = Vec::new();

// Преобразование весов и смещений в параметры для оптимизации

for i in 0..self.output_size {

for j in 0..self.input_size {

let param = crate::autodiff::Parameter::new(

self.weights[i][j].into(),

format!("weight_{}_{}", i, j)

);

params.push(param);

}

let bias_param = crate::autodiff::Parameter::new(

self.biases[i].into(),

format!("bias_{}", i)

);

params.push(bias_param);

}

params

}

}

// Блочная обработка для больших входных данных

impl<T> LinearLayer<T>

where

T: Copy + Default + Add<Output = T> + Mul<Output = T> + From<f32>,

{

pub fn process_batch(&self, inputs: &[Vec<T>]) -> Result<Vec<Vec<T>>, LinearLayerError> {

let batch_size = inputs.len();

let mut results = Vec::with_capacity(batch_size);

for (batch_idx, input) in inputs.iter().enumerate() {

if input.len() != self.input_size {

return Err(LinearLayerError::InputSizeMismatch {

input_size: input.len(),

expected: self.input_size,

});

}

let mut output = Vec::with_capacity(self.output_size);

for i in 0..self.output_size {

let mut sum = T::default();

for j in 0..self.input_size {

sum = sum + input[j] * self.weights[i][j];

}

output.push(sum + self.biases[i]);

}

results.push(output);

}

Ok(results)

}

}

// Специализированная реализация для f32 с дополнительными оптимизациями

impl LinearLayer<f32> {

pub fn quantize_weights(&mut self, levels: usize) -> Result<(), LinearLayerError> {

if levels < 2 {

return Err(LinearLayerError::InitializationError {

reason: "Уровней квантования должно быть не менее 2".to_string(),

});

}

// Находим минимальное и максимальное значения весов

let (min, max) = self.weights.iter()

.flatten()

.fold((f32::INFINITY, f32::NEG_INFINITY), |(min, max), &w| {

(min.min(w), max.max(w))

});

let step = (max - min) / (levels - 1) as f32;

// Квантуем веса

for i in 0..self.output_size {

for j in 0..self.input_size {

let quantized = ((self.weights[i][j] - min) / step).round() * step + min;

self.weights[i][j] = quantized;

self.weight_cache.insert((i, j), quantized);

}

}

Ok(())

}

}

// Тестирование реализации

#[cfg(test)]

mod tests {

use super::*;

#[test]

fn test_linear_layer_creation() {

let layer = LinearLayer::<f32>::new(10, 5).unwrap();

assert_eq!(layer.input_size, 10);

assert_eq!(layer.output_size, 5);

}

#[test]

fn test_linear_layer_forward() {

let mut layer = LinearLayer::<f32>::new(3, 2).unwrap();

// Устанавливаем конкретные веса для тестирования

layer.set_weight(0, 0, 0.1).unwrap();

layer.set_weight(0, 1, 0.2).unwrap();

layer.set_weight(0, 2, 0.3).unwrap();

layer.set_weight(1, 0, 0.4).unwrap();

layer.set_weight(1, 1, 0.5).unwrap();

layer.set_weight(1, 2, 0.6).unwrap();

layer.biases = vec![0.1, 0.2];

let input = vec![1.0, 2.0, 3.0];

let output = layer.forward(&input).unwrap();

// Проверяем вычисления вручную

let expected0 = 1.0*0.1 + 2.0*0.2 + 3.0*0.3 + 0.1;

let expected1 = 1.0*0.4 + 2.0*0.5 + 3.0*0.6 + 0.2;

assert!((output[0] - expected0).abs() < 1e-6);

assert!((output[1] - expected1).abs() < 1e-6);

}

#[test]

fn test_error_handling() {

let layer = LinearLayer::<f32>::new(0, 5);

assert!(layer.is_err());

let mut layer = LinearLayer::<f32>::new(3, 2).unwrap();

let result = layer.forward(&vec![1.0, 2.0]);

assert!(result.is_err());

}

}