rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
redclawsystems
2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/models/rtx-vision/src/models/configs.rs
+111
View File
@@ -0,0 +1,111 @@
//! Model configurations for vision models
/// Vision Transformer configuration
#[derive(Clone, Debug)]
pub struct ViTConfig {
pub image_size: usize,
pub patch_size: usize,
pub num_classes: usize,
pub embed_dim: usize,
pub depth: usize,
pub num_heads: usize,
pub mlp_ratio: f32,
pub dropout: f32,
pub attention_dropout: f32,
}
impl ViTConfig {
/// ViT-Base/16 configuration
pub fn base_16() -> Self {
Self {
image_size: 224,
patch_size: 16,
num_classes: 1000,
embed_dim: 768,
depth: 12,
num_heads: 12,
mlp_ratio: 4.0,
dropout: 0.0,
attention_dropout: 0.0,
}
}
/// ViT-Large/16 configuration
pub fn large_16() -> Self {
Self {
image_size: 224,
patch_size: 16,
num_classes: 1000,
embed_dim: 1024,
depth: 24,
num_heads: 16,
mlp_ratio: 4.0,
dropout: 0.0,
attention_dropout: 0.0,
}
}
/// ViT-Huge/14 configuration
pub fn huge_14() -> Self {
Self {
image_size: 224,
patch_size: 14,
num_classes: 1000,
embed_dim: 1280,
depth: 32,
num_heads: 16,
mlp_ratio: 4.0,
dropout: 0.0,
attention_dropout: 0.0,
}
}
}
/// ConvNeXt configuration
#[derive(Clone, Debug)]
pub struct ConvNeXtConfig {
pub in_channels: usize,
pub num_classes: usize,
pub depths: Vec<usize>,
pub dims: Vec<usize>,
pub drop_path_rate: f32,
pub layer_scale_init_value: f32,
}
impl ConvNeXtConfig {
/// ConvNeXt-Tiny configuration
pub fn tiny() -> Self {
Self {
in_channels: 3,
num_classes: 1000,
depths: vec![3, 3, 9, 3],
dims: vec![96, 192, 384, 768],
drop_path_rate: 0.1,
layer_scale_init_value: 1e-6,
}
}
/// ConvNeXt-Small configuration
pub fn small() -> Self {
Self {
in_channels: 3,
num_classes: 1000,
depths: vec![3, 3, 27, 3],
dims: vec![96, 192, 384, 768],
drop_path_rate: 0.4,
layer_scale_init_value: 1e-6,
}
}
/// ConvNeXt-Base configuration
pub fn base() -> Self {
Self {
in_channels: 3,
num_classes: 1000,
depths: vec![3, 3, 27, 3],
dims: vec![128, 256, 512, 1024],
drop_path_rate: 0.5,
layer_scale_init_value: 1e-6,
}
}
}
crates/models/rtx-vision/src/models/convnext.rs
+117
View File
@@ -0,0 +1,117 @@
//! ConvNeXt implementation - modern ConvNet architecture
use crate::error::{Result, VisionError};
use crate::models::configs::ConvNeXtConfig;
use crate::{Device, Tensor};
/// ConvNeXt block
#[derive(Clone, Debug)]
pub struct ConvNeXtBlock {
dim: usize,
drop_path: f32,
layer_scale: Option<Tensor>,
device: Device,
}
impl ConvNeXtBlock {
pub fn new(dim: usize, drop_path: f32, layer_scale_init: f32, device: &Device) -> Result<Self> {
let layer_scale = if layer_scale_init > 0.0 {
Some(Tensor::ones([dim], device)?.mul_scalar(layer_scale_init)?)
} else {
None
};
Ok(Self {
dim,
drop_path,
layer_scale,
device: device.clone(),
})
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
// Simplified forward pass
// In production, would implement depthwise conv, LayerNorm, pointwise convs
Ok(x.clone())
}
}
/// ConvNeXt model
#[derive(Clone, Debug)]
pub struct ConvNeXt {
config: ConvNeXtConfig,
stages: Vec<Vec<ConvNeXtBlock>>,
head: Tensor,
device: Device,
}
impl ConvNeXt {
/// Create a new ConvNeXt model
pub fn new(config: ConvNeXtConfig, device: &Device) -> Result<Self> {
let mut stages = Vec::new();
for (stage_idx, &depth) in config.depths.iter().enumerate() {
let dim = config.dims[stage_idx];
let mut blocks = Vec::new();
for _ in 0..depth {
blocks.push(ConvNeXtBlock::new(
dim,
config.drop_path_rate,
config.layer_scale_init_value,
device,
)?);
}
stages.push(blocks);
}
// Classification head
let final_dim = *config.dims.last().unwrap();
let head = Tensor::randn(&[final_dim, config.num_classes], device)?;
Ok(Self {
config,
stages,
head,
device: device.clone(),
})
}
/// Forward pass through ConvNeXt
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
let mut x = x.clone();
// Process through stages
for stage in &self.stages {
for block in stage {
x = block.forward(&x)?;
}
}
// Global average pooling
x = self.global_avg_pool(&x)?;
// Classification head
x.matmul(&self.head.transpose(-2, -1)?)
.map_err(VisionError::from)
}
fn global_avg_pool(&self, x: &Tensor) -> Result<Tensor> {
// Simplified global average pooling
// In production, would properly average over spatial dimensions
let shape = x.shape().dims();
match shape.len() {
4 => {
// [B, C, H, W] -> [B, C]
let batch = shape[0];
let channels = shape[1];
Tensor::randn(&[batch, channels], &self.device).map_err(VisionError::from)
}
_ => Err(VisionError::InvalidDimensions {
expected: "4D tensor".to_string(),
got: format!("{:?}", shape),
}),
}
}
}
crates/models/rtx-vision/src/models/mod.rs
+9
View File
@@ -0,0 +1,9 @@
//! Vision model architectures
mod configs;
mod convnext;
mod vit;
pub use configs::{ConvNeXtConfig, ViTConfig};
pub use convnext::{ConvNeXt, ConvNeXtBlock};
pub use vit::{ViT, ViTBlock};
crates/models/rtx-vision/src/models/vit.rs
+134
View File
@@ -0,0 +1,134 @@
//! Vision Transformer (ViT) implementation
use crate::error::{Result, VisionError};
use crate::layers::{LayerNorm, PatchEmbedding, PositionalEncodingType};
use crate::models::configs::ViTConfig;
use crate::{Device, Tensor};
/// Vision Transformer block
#[derive(Clone, Debug)]
pub struct ViTBlock {
embed_dim: usize,
num_heads: usize,
mlp_ratio: f32,
norm1: LayerNorm,
norm2: LayerNorm,
device: Device,
}
impl ViTBlock {
pub fn new(
embed_dim: usize,
num_heads: usize,
mlp_ratio: f32,
device: &Device,
) -> Result<Self> {
let norm1 = LayerNorm::new(vec![embed_dim], 1e-6, device)?;
let norm2 = LayerNorm::new(vec![embed_dim], 1e-6, device)?;
Ok(Self {
embed_dim,
num_heads,
mlp_ratio,
norm1,
norm2,
device: device.clone(),
})
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
// Simplified forward pass
// In production, would implement full attention and MLP
Ok(x.clone())
}
}
/// Vision Transformer model
#[derive(Clone, Debug)]
pub struct ViT {
config: ViTConfig,
patch_embed: PatchEmbedding,
blocks: Vec<ViTBlock>,
norm: LayerNorm,
head: Tensor,
device: Device,
}
impl ViT {
/// Create a new Vision Transformer
pub fn new(config: ViTConfig, device: &Device) -> Result<Self> {
let patch_embed = PatchEmbedding::new(
config.image_size,
config.patch_size,
3,
config.embed_dim,
device,
)?
.with_positional_encoding(PositionalEncodingType::Learnable)
.with_class_token();
let blocks: Result<Vec<_>> = (0..config.depth)
.map(|_| ViTBlock::new(config.embed_dim, config.num_heads, config.mlp_ratio, device))
.collect();
let blocks = blocks?;
let norm = LayerNorm::new(vec![config.embed_dim], 1e-6, device)?;
// Classification head
let head = Tensor::randn(&[config.embed_dim, config.num_classes], device)?;
Ok(Self {
config,
patch_embed,
blocks,
norm,
head,
device: device.clone(),
})
}
/// Forward pass through ViT
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
// Patch embedding
let mut x = self.patch_embed.forward(x)?;
// Add class token
x = self.patch_embed.add_class_token(&x)?;
// Transformer blocks
for block in &self.blocks {
x = block.forward(&x)?;
}
// Layer norm
x = self.norm.forward(&x)?;
// Extract class token (first token)
let cls_token = self.extract_class_token(&x)?;
// Classification head
cls_token
.matmul(&self.head.transpose(-2, -1)?)
.map_err(VisionError::from)
}
fn extract_class_token(&self, x: &Tensor) -> Result<Tensor> {
let shape = x.shape().dims();
match shape.len() {
2 => {
// [seq_len, embed_dim] -> [1, embed_dim]
x.narrow(0, 0, 1).map_err(VisionError::from)
}
3 => {
// [batch, seq_len, embed_dim] -> [batch, embed_dim]
x.narrow(1, 0, 1)?
.squeeze(Some(1))
.map_err(VisionError::from)
}
_ => Err(VisionError::InvalidDimensions {
expected: "2D or 3D tensor".to_string(),
got: format!("{:?}", shape),
}),
}
}
}