298 lines
8.1 KiB
Rust
298 lines
8.1 KiB
Rust
//! Tests for patch embedding layer
|
|
//!
|
|
//! TDD: Define expected behavior for Vision Transformer patch embedding
|
|
|
|
use approx::assert_abs_diff_eq;
|
|
use rtx_tensor::{Device, Tensor};
|
|
use rtx_vision::Result;
|
|
use rtx_vision::layers::{PatchEmbedding, PositionalEncoding, PositionalEncodingType};
|
|
use rtx_vision::preprocessing::ImageTensor;
|
|
|
|
#[test]
|
|
fn test_patch_embedding_creation() {
|
|
let device = Device::cpu();
|
|
|
|
// Standard ViT-B/16 configuration
|
|
let patch_embed = PatchEmbedding::new(
|
|
224, // image_size
|
|
16, // patch_size
|
|
3, // in_channels
|
|
768, // embed_dim
|
|
&device,
|
|
);
|
|
|
|
assert!(patch_embed.is_ok());
|
|
let patch_embed = patch_embed.unwrap();
|
|
|
|
// Should create (224/16)^2 = 196 patches
|
|
assert_eq!(patch_embed.num_patches(), 196);
|
|
assert_eq!(patch_embed.embed_dim(), 768);
|
|
}
|
|
|
|
#[test]
|
|
fn test_patch_embedding_forward() {
|
|
let device = Device::cpu();
|
|
|
|
// Create patch embedding layer
|
|
let patch_embed = PatchEmbedding::new(
|
|
224, // image_size
|
|
16, // patch_size
|
|
3, // in_channels
|
|
768, // embed_dim
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Create test image [3, 224, 224]
|
|
let image_data = vec![0.5f32; 3 * 224 * 224];
|
|
let image = ImageTensor::from_array(image_data, 224, 224, 3, &device).unwrap();
|
|
|
|
// Forward pass
|
|
let patches = patch_embed.forward(image.to_tensor());
|
|
assert!(patches.is_ok());
|
|
|
|
let patches = patches.unwrap();
|
|
// Output shape should be [196, 768] (num_patches, embed_dim)
|
|
assert_eq!(patches.shape().dims(), &[196, 768]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_patch_embedding_with_batch() {
|
|
let device = Device::cpu();
|
|
|
|
let patch_embed = PatchEmbedding::new(
|
|
224, // image_size
|
|
16, // patch_size
|
|
3, // in_channels
|
|
768, // embed_dim
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Create batch of images [B, C, H, W]
|
|
let batch_tensor = Tensor::randn(&[4, 3, 224, 224], &device).unwrap();
|
|
|
|
let patches = patch_embed.forward(&batch_tensor);
|
|
assert!(patches.is_ok());
|
|
|
|
let patches = patches.unwrap();
|
|
// Output shape should be [4, 196, 768] (batch, num_patches, embed_dim)
|
|
assert_eq!(patches.shape().dims(), &[4, 196, 768]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_patch_embedding_different_sizes() {
|
|
let device = Device::cpu();
|
|
|
|
// Test with 32x32 patches (used in some ViT variants)
|
|
let patch_embed = PatchEmbedding::new(
|
|
224, // image_size
|
|
32, // patch_size
|
|
3, // in_channels
|
|
768, // embed_dim
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Should create (224/32)^2 = 49 patches
|
|
assert_eq!(patch_embed.num_patches(), 49);
|
|
|
|
// Test with non-square image
|
|
let patch_embed_rect = PatchEmbedding::new_rectangular(
|
|
256, // height
|
|
128, // width
|
|
16, // patch_size
|
|
3, // in_channels
|
|
768, // embed_dim
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Should create (256/16) * (128/16) = 16 * 8 = 128 patches
|
|
assert_eq!(patch_embed_rect.num_patches(), 128);
|
|
}
|
|
|
|
#[test]
|
|
fn test_class_token() {
|
|
let device = Device::cpu();
|
|
|
|
// Create patch embedding with class token
|
|
let patch_embed = PatchEmbedding::new(
|
|
224, // image_size
|
|
16, // patch_size
|
|
3, // in_channels
|
|
768, // embed_dim
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Create test image
|
|
let image = Tensor::randn(&[3, 224, 224], &device).unwrap();
|
|
let patches = patch_embed.forward(&image).unwrap();
|
|
|
|
// Add class token
|
|
let patches_with_cls = patch_embed.add_class_token(&patches);
|
|
assert!(patches_with_cls.is_ok());
|
|
|
|
let patches_with_cls = patches_with_cls.unwrap();
|
|
// Should have one additional token: [197, 768] instead of [196, 768]
|
|
assert_eq!(patches_with_cls.shape().dims(), &[197, 768]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_learnable_positional_encoding() {
|
|
let device = Device::cpu();
|
|
|
|
// Create learnable positional encoding
|
|
let pos_enc = PositionalEncoding::new(
|
|
196, // num_patches (14x14 for 224x224 with patch_size=16)
|
|
768, // embed_dim
|
|
PositionalEncodingType::Learnable,
|
|
&device,
|
|
);
|
|
|
|
assert!(pos_enc.is_ok());
|
|
let pos_enc = pos_enc.unwrap();
|
|
|
|
// Create patch embeddings
|
|
let patches = Tensor::randn(&[196, 768], &device).unwrap();
|
|
|
|
// Add positional encoding
|
|
let encoded = pos_enc.forward(&patches);
|
|
assert!(encoded.is_ok());
|
|
|
|
let encoded = encoded.unwrap();
|
|
assert_eq!(encoded.shape().dims(), &[196, 768]);
|
|
|
|
// Values should be different from input (positional info added)
|
|
let patches_data = patches.to_vec().unwrap();
|
|
let encoded_data = encoded.to_vec().unwrap();
|
|
|
|
let changed = patches_data
|
|
.iter()
|
|
.zip(encoded_data.iter())
|
|
.any(|(a, b)| (a - b).abs() > 1e-5);
|
|
assert!(changed, "Positional encoding should modify embeddings");
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Sinusoidal positional encoding implementation differs"]
|
|
fn test_sinusoidal_positional_encoding() {
|
|
let device = Device::cpu();
|
|
|
|
// Create sinusoidal positional encoding (no learnable parameters)
|
|
let pos_enc = PositionalEncoding::new(
|
|
196, // num_patches
|
|
768, // embed_dim
|
|
PositionalEncodingType::Sinusoidal,
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Create patch embeddings
|
|
let patches = Tensor::zeros(&[196, 768], &device).unwrap();
|
|
|
|
// Add positional encoding
|
|
let encoded = pos_enc.forward(&patches).unwrap();
|
|
|
|
// Check that sinusoidal pattern is applied
|
|
let encoded_data = encoded.to_vec().unwrap();
|
|
|
|
// First position should have specific sinusoidal values
|
|
// sin(0) = 0, cos(0) = 1 for even dimensions
|
|
assert_abs_diff_eq!(encoded_data[0], 0.0, epsilon = 1e-5); // sin(0)
|
|
assert_abs_diff_eq!(encoded_data[1], 1.0, epsilon = 1e-5); // cos(0)
|
|
}
|
|
|
|
#[test]
|
|
fn test_patch_embedding_with_positional_encoding() {
|
|
let device = Device::cpu();
|
|
|
|
// Create complete patch embedding with positional encoding
|
|
let patch_embed = PatchEmbedding::new(
|
|
224, // image_size
|
|
16, // patch_size
|
|
3, // in_channels
|
|
768, // embed_dim
|
|
&device,
|
|
)
|
|
.unwrap()
|
|
.with_positional_encoding(PositionalEncodingType::Learnable);
|
|
|
|
// Create test image
|
|
let image = Tensor::randn(&[3, 224, 224], &device).unwrap();
|
|
|
|
// Forward pass should include positional encoding
|
|
let output = patch_embed.forward(&image).unwrap();
|
|
assert_eq!(output.shape().dims(), &[196, 768]);
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires_grad propagation incomplete"]
|
|
fn test_patch_embedding_gradients() {
|
|
let device = Device::cpu();
|
|
|
|
// Create patch embedding layer
|
|
let patch_embed = PatchEmbedding::new(
|
|
32, // smaller image for faster test
|
|
8, // patch_size
|
|
3, // in_channels
|
|
256, // embed_dim
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Create test image with requires_grad
|
|
let mut image = Tensor::randn(&[3, 32, 32], &device).unwrap();
|
|
image.set_requires_grad(true);
|
|
|
|
// Forward pass
|
|
let patches = patch_embed.forward(&image).unwrap();
|
|
|
|
// Simulate loss and backward
|
|
let loss = patches.sum(None).unwrap();
|
|
|
|
// Check that gradients can flow
|
|
assert!(patches.requires_grad());
|
|
|
|
// In real implementation, would check actual gradient values
|
|
}
|
|
|
|
#[test]
|
|
fn test_invalid_patch_size() {
|
|
let device = Device::cpu();
|
|
|
|
// Image size not divisible by patch size
|
|
let result = PatchEmbedding::new(
|
|
224, // image_size
|
|
15, // patch_size (224 not divisible by 15)
|
|
3, // in_channels
|
|
768, // embed_dim
|
|
&device,
|
|
);
|
|
|
|
assert!(result.is_err());
|
|
assert!(result.unwrap_err().to_string().contains("patch"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_patch_embedding_2d_conv() {
|
|
let device = Device::cpu();
|
|
|
|
// Patch embedding uses Conv2d with kernel_size=patch_size, stride=patch_size
|
|
let patch_embed = PatchEmbedding::new(
|
|
224, // image_size
|
|
16, // patch_size
|
|
3, // in_channels
|
|
768, // embed_dim
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Verify internal conv parameters
|
|
assert_eq!(patch_embed.kernel_size(), 16);
|
|
assert_eq!(patch_embed.stride(), 16);
|
|
assert_eq!(patch_embed.in_channels(), 3);
|
|
assert_eq!(patch_embed.out_channels(), 768);
|
|
}
|