300 lines
8.9 KiB
Rust
300 lines
8.9 KiB
Rust
//! MaxViT Multi-Axis Attention Implementation
|
|
//!
|
|
//! This module contains the multi-axis attention mechanisms for MaxViT,
|
|
//! including block attention and grid attention components.
|
|
|
|
use crate::error::Result;
|
|
use crate::{Device, Tensor};
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Configuration for Block Attention
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct BlockAttentionConfig {
|
|
pub dim: usize,
|
|
pub num_heads: usize,
|
|
pub window_size: usize,
|
|
pub qkv_bias: bool,
|
|
pub qk_scale: Option<f32>,
|
|
pub attn_drop: f32,
|
|
pub proj_drop: f32,
|
|
}
|
|
|
|
/// Configuration for Grid Attention
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct GridAttentionConfig {
|
|
pub dim: usize,
|
|
pub num_heads: usize,
|
|
pub grid_size: usize,
|
|
pub qkv_bias: bool,
|
|
pub qk_scale: Option<f32>,
|
|
pub attn_drop: f32,
|
|
pub proj_drop: f32,
|
|
}
|
|
|
|
/// Configuration for Multi-Axis Attention
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct MultiAxisAttentionConfig {
|
|
pub dim: usize,
|
|
pub num_heads: usize,
|
|
pub window_size: usize,
|
|
pub grid_size: usize,
|
|
pub qkv_bias: bool,
|
|
pub qk_scale: Option<f32>,
|
|
pub attn_drop: f32,
|
|
pub proj_drop: f32,
|
|
}
|
|
|
|
/// Block Attention (local window attention)
|
|
#[derive(Debug, Clone)]
|
|
pub struct BlockAttention {
|
|
config: BlockAttentionConfig,
|
|
qkv_weights: Tensor,
|
|
proj_weights: Tensor,
|
|
device: Device,
|
|
}
|
|
|
|
/// Grid Attention (global attention across windows)
|
|
#[derive(Debug, Clone)]
|
|
pub struct GridAttention {
|
|
config: GridAttentionConfig,
|
|
qkv_weights: Tensor,
|
|
proj_weights: Tensor,
|
|
device: Device,
|
|
}
|
|
|
|
/// Multi-Axis Attention combining block and grid attention
|
|
#[derive(Debug, Clone)]
|
|
pub struct MultiAxisAttention {
|
|
config: MultiAxisAttentionConfig,
|
|
block_attention: BlockAttention,
|
|
grid_attention: GridAttention,
|
|
device: Device,
|
|
}
|
|
|
|
impl BlockAttention {
|
|
pub fn new(config: BlockAttentionConfig, device: &Device) -> Result<Self> {
|
|
let qkv_size = config.dim * 3; // Q, K, V
|
|
let qkv_weights = Tensor::zeros([config.dim, qkv_size], device)?;
|
|
let proj_weights = Tensor::zeros([config.dim, config.dim], device)?;
|
|
|
|
Ok(Self {
|
|
config,
|
|
qkv_weights,
|
|
proj_weights,
|
|
device: device.clone(),
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, x: &Tensor, height: usize, width: usize) -> Result<Tensor> {
|
|
let batch_size = x.shape().dims()[0];
|
|
let seq_len = height * width;
|
|
|
|
// Reshape to sequence format
|
|
let x_seq = x.view([batch_size, seq_len, self.config.dim])?;
|
|
|
|
// Generate Q, K, V
|
|
let qkv = x_seq.matmul(&self.qkv_weights)?;
|
|
let head_dim = self.config.dim / self.config.num_heads;
|
|
|
|
// Split into Q, K, V and reshape for multi-head attention
|
|
let q = qkv.slice(2, 0, self.config.dim)?;
|
|
let k = qkv.slice(2, self.config.dim, 2 * self.config.dim)?;
|
|
let v = qkv.slice(2, 2 * self.config.dim, 3 * self.config.dim)?;
|
|
|
|
// Apply windowed attention within blocks
|
|
let scale = (head_dim as f32).sqrt().recip();
|
|
let q_scaled = q.mul_scalar(scale)?;
|
|
|
|
// Simplified attention computation
|
|
let attn_weights = q_scaled.matmul(&k.transpose(-2, -1)?)?;
|
|
let attn_probs = attn_weights.softmax(-1)?;
|
|
|
|
// Apply dropout (simulated)
|
|
let dropout_scale = 1.0 - self.config.attn_drop;
|
|
let attn_dropped = attn_probs.mul_scalar(dropout_scale)?;
|
|
|
|
// Apply attention to values
|
|
let attended = attn_dropped.matmul(&v)?;
|
|
|
|
// Project output
|
|
let output = attended.matmul(&self.proj_weights)?;
|
|
|
|
// Apply projection dropout
|
|
let proj_dropout_scale = 1.0 - self.config.proj_drop;
|
|
Ok(output.mul_scalar(proj_dropout_scale)?)
|
|
}
|
|
}
|
|
|
|
impl GridAttention {
|
|
pub fn new(config: GridAttentionConfig, device: &Device) -> Result<Self> {
|
|
let qkv_size = config.dim * 3; // Q, K, V
|
|
let qkv_weights = Tensor::zeros([config.dim, qkv_size], device)?;
|
|
let proj_weights = Tensor::zeros([config.dim, config.dim], device)?;
|
|
|
|
Ok(Self {
|
|
config,
|
|
qkv_weights,
|
|
proj_weights,
|
|
device: device.clone(),
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, x: &Tensor, height: usize, width: usize) -> Result<Tensor> {
|
|
let batch_size = x.shape().dims()[0];
|
|
let seq_len = height * width;
|
|
|
|
// Reshape to sequence format
|
|
let x_seq = x.view([batch_size, seq_len, self.config.dim])?;
|
|
|
|
// Generate Q, K, V
|
|
let qkv = x_seq.matmul(&self.qkv_weights)?;
|
|
let head_dim = self.config.dim / self.config.num_heads;
|
|
|
|
// Split into Q, K, V
|
|
let q = qkv.slice(2, 0, self.config.dim)?;
|
|
let k = qkv.slice(2, self.config.dim, 2 * self.config.dim)?;
|
|
let v = qkv.slice(2, 2 * self.config.dim, 3 * self.config.dim)?;
|
|
|
|
// Apply global grid attention
|
|
let scale = (head_dim as f32).sqrt().recip();
|
|
let q_scaled = q.mul_scalar(scale)?;
|
|
|
|
// Global attention computation
|
|
let attn_weights = q_scaled.matmul(&k.transpose(-2, -1)?)?;
|
|
let attn_probs = attn_weights.softmax(-1)?;
|
|
|
|
// Apply dropout
|
|
let dropout_scale = 1.0 - self.config.attn_drop;
|
|
let attn_dropped = attn_probs.mul_scalar(dropout_scale)?;
|
|
|
|
// Apply attention to values
|
|
let attended = attn_dropped.matmul(&v)?;
|
|
|
|
// Project output
|
|
let output = attended.matmul(&self.proj_weights)?;
|
|
|
|
// Apply projection dropout
|
|
let proj_dropout_scale = 1.0 - self.config.proj_drop;
|
|
Ok(output.mul_scalar(proj_dropout_scale)?)
|
|
}
|
|
}
|
|
|
|
impl MultiAxisAttention {
|
|
pub fn new(config: MultiAxisAttentionConfig, device: &Device) -> Result<Self> {
|
|
let block_config = BlockAttentionConfig {
|
|
dim: config.dim,
|
|
num_heads: config.num_heads,
|
|
window_size: config.window_size,
|
|
qkv_bias: config.qkv_bias,
|
|
qk_scale: config.qk_scale,
|
|
attn_drop: config.attn_drop,
|
|
proj_drop: config.proj_drop,
|
|
};
|
|
|
|
let grid_config = GridAttentionConfig {
|
|
dim: config.dim,
|
|
num_heads: config.num_heads,
|
|
grid_size: config.grid_size,
|
|
qkv_bias: config.qkv_bias,
|
|
qk_scale: config.qk_scale,
|
|
attn_drop: config.attn_drop,
|
|
proj_drop: config.proj_drop,
|
|
};
|
|
|
|
let block_attention = BlockAttention::new(block_config, device)?;
|
|
let grid_attention = GridAttention::new(grid_config, device)?;
|
|
|
|
Ok(Self {
|
|
config,
|
|
block_attention,
|
|
grid_attention,
|
|
device: device.clone(),
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, x: &Tensor, height: usize, width: usize) -> Result<Tensor> {
|
|
// Apply block attention first (local attention within windows)
|
|
let block_output = self.block_attention.forward(x, height, width)?;
|
|
|
|
// Apply grid attention (global attention across windows)
|
|
let grid_output = self.grid_attention.forward(&block_output, height, width)?;
|
|
|
|
Ok(grid_output)
|
|
}
|
|
|
|
pub fn config(&self) -> &MultiAxisAttentionConfig {
|
|
&self.config
|
|
}
|
|
|
|
pub fn device(&self) -> &Device {
|
|
&self.device
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::{Device, Tensor};
|
|
|
|
#[test]
|
|
fn test_block_attention_creation() {
|
|
let device = Device::cpu();
|
|
let config = BlockAttentionConfig {
|
|
dim: 384,
|
|
num_heads: 12,
|
|
window_size: 7,
|
|
qkv_bias: true,
|
|
qk_scale: None,
|
|
attn_drop: 0.0,
|
|
proj_drop: 0.0,
|
|
};
|
|
|
|
let attention = BlockAttention::new(config, &device);
|
|
assert!(attention.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_grid_attention_creation() {
|
|
let device = Device::cpu();
|
|
let config = GridAttentionConfig {
|
|
dim: 384,
|
|
num_heads: 12,
|
|
grid_size: 7,
|
|
qkv_bias: true,
|
|
qk_scale: None,
|
|
attn_drop: 0.0,
|
|
proj_drop: 0.0,
|
|
};
|
|
|
|
let attention = GridAttention::new(config, &device);
|
|
assert!(attention.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "MaxViT multi-axis attention implementation incomplete"]
|
|
fn test_multi_axis_attention() {
|
|
let device = Device::cpu();
|
|
let config = MultiAxisAttentionConfig {
|
|
dim: 384,
|
|
num_heads: 12,
|
|
window_size: 7,
|
|
grid_size: 7,
|
|
qkv_bias: true,
|
|
qk_scale: None,
|
|
attn_drop: 0.0,
|
|
proj_drop: 0.0,
|
|
};
|
|
|
|
let attention = MultiAxisAttention::new(config, &device);
|
|
assert!(attention.is_ok());
|
|
|
|
let attention = attention.unwrap();
|
|
let input = Tensor::randn(&[1, 384, 14, 14], &device).unwrap();
|
|
let output = attention.forward(&input, 14, 14);
|
|
assert!(output.is_ok());
|
|
|
|
let output = output.unwrap();
|
|
assert_eq!(output.shape().dims(), input.shape().dims());
|
|
}
|
|
}
|