Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,354 @@
|
||||
use crate::Graph;
|
||||
use crate::error::{GeomError, Result};
|
||||
use crate::layers::GNNLayer;
|
||||
use rtx_tensor::Tensor;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Graph Attention Network (GAT) Layer
|
||||
///
|
||||
/// Implements multi-head attention mechanism from Veličković et al. (2018):
|
||||
/// α_ij = softmax(LeakyReLU(a^T [W h_i || W h_j]))
|
||||
/// h_i' = σ(Σ_j α_ij W h_j)
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GATLayer {
|
||||
weight: Tensor,
|
||||
attention_weight: Tensor,
|
||||
bias: Tensor,
|
||||
input_dim: usize,
|
||||
output_dim: usize,
|
||||
num_heads: usize,
|
||||
gradients_enabled: bool,
|
||||
computed_gradients: bool,
|
||||
attention_weights: HashMap<(usize, usize), f32>, // (from, to) -> attention weight
|
||||
}
|
||||
|
||||
impl GATLayer {
|
||||
pub fn new(input_dim: usize, output_dim: usize, num_heads: usize) -> Result<Self> {
|
||||
if input_dim == 0 || output_dim == 0 || num_heads == 0 {
|
||||
return Err(GeomError::LayerInitializationFailed {
|
||||
reason: "Dimensions and num_heads must be greater than 0".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
if !output_dim.is_multiple_of(num_heads) {
|
||||
return Err(GeomError::LayerInitializationFailed {
|
||||
reason: "Output dimension must be divisible by number of heads".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
let head_dim = output_dim / num_heads;
|
||||
|
||||
// Xavier initialization
|
||||
let fan_in = input_dim as f32;
|
||||
let fan_out = output_dim as f32;
|
||||
let limit = (6.0 / (fan_in + fan_out)).sqrt();
|
||||
|
||||
// Weight matrix for feature transformation
|
||||
let mut weight_data = Vec::with_capacity(input_dim * output_dim);
|
||||
for i in 0..(input_dim * output_dim) {
|
||||
let val = (((i * 7919 + 1337) % 10000) as f32 / 5000.0 - 1.0) * limit;
|
||||
weight_data.push(val);
|
||||
}
|
||||
let weight = Tensor::from_data(
|
||||
weight_data,
|
||||
[input_dim, output_dim],
|
||||
&rtx_tensor::Device::default(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Attention weight vector (2 * head_dim for concatenated features)
|
||||
let attention_dim = 2 * head_dim * num_heads;
|
||||
let mut attention_data = Vec::with_capacity(attention_dim);
|
||||
for i in 0..attention_dim {
|
||||
let val = (((i * 6151 + 2357) % 10000) as f32 / 5000.0 - 1.0) * limit;
|
||||
attention_data.push(val);
|
||||
}
|
||||
let attention_weight = Tensor::from_data(
|
||||
attention_data,
|
||||
[attention_dim],
|
||||
&rtx_tensor::Device::default(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// Bias
|
||||
let bias = Tensor::zeros([output_dim], &rtx_tensor::Device::default()).unwrap();
|
||||
|
||||
Ok(Self {
|
||||
weight,
|
||||
attention_weight,
|
||||
bias,
|
||||
input_dim,
|
||||
output_dim,
|
||||
num_heads,
|
||||
gradients_enabled: false,
|
||||
computed_gradients: false,
|
||||
attention_weights: HashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_attention_weights(&self) -> &HashMap<(usize, usize), f32> {
|
||||
&self.attention_weights
|
||||
}
|
||||
|
||||
fn compute_attention(&mut self, graph: &Graph) -> Result<HashMap<usize, Vec<(usize, f32)>>> {
|
||||
let mut attention_map = HashMap::new();
|
||||
|
||||
for node_id in graph.nodes() {
|
||||
let neighbors = graph.incoming_neighbors(node_id);
|
||||
let mut node_attentions = Vec::new();
|
||||
|
||||
if !neighbors.is_empty() {
|
||||
let node_features = graph.node(node_id).unwrap().features();
|
||||
|
||||
let _attention_scores: Vec<f32> = Vec::new();
|
||||
let mut raw_scores = Vec::new();
|
||||
|
||||
// Compute raw attention scores for all neighbors
|
||||
for &neighbor_id in &neighbors {
|
||||
let neighbor_features = graph.node(neighbor_id).unwrap().features();
|
||||
|
||||
// Compute attention score (simplified)
|
||||
let node_2d = node_features.reshape([1, self.input_dim]).map_err(|e| {
|
||||
GeomError::ForwardPassFailed {
|
||||
reason: format!("Node feature reshape failed: {e}"),
|
||||
}
|
||||
})?;
|
||||
|
||||
let neighbor_2d =
|
||||
neighbor_features
|
||||
.reshape([1, self.input_dim])
|
||||
.map_err(|e| GeomError::ForwardPassFailed {
|
||||
reason: format!("Neighbor feature reshape failed: {e}"),
|
||||
})?;
|
||||
|
||||
let node_transformed =
|
||||
node_2d
|
||||
.matmul(&self.weight)
|
||||
.map_err(|e| GeomError::ForwardPassFailed {
|
||||
reason: format!("Node transformation failed: {e}"),
|
||||
})?;
|
||||
|
||||
let neighbor_transformed = neighbor_2d.matmul(&self.weight).map_err(|e| {
|
||||
GeomError::ForwardPassFailed {
|
||||
reason: format!("Neighbor transformation failed: {e}"),
|
||||
}
|
||||
})?;
|
||||
|
||||
// Simplified attention computation: dot product of first elements
|
||||
let node_data = node_transformed.to_cpu().unwrap();
|
||||
let neighbor_data = neighbor_transformed.to_cpu().unwrap();
|
||||
|
||||
let raw_score = if !node_data.is_empty() && !neighbor_data.is_empty() {
|
||||
node_data[0] * neighbor_data[0] // Simplified attention mechanism
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
raw_scores.push(raw_score);
|
||||
}
|
||||
|
||||
// Apply softmax to get attention weights
|
||||
let max_score = raw_scores.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
||||
let exp_scores: Vec<f32> =
|
||||
raw_scores.iter().map(|&s| (s - max_score).exp()).collect();
|
||||
let sum_exp = exp_scores.iter().sum::<f32>();
|
||||
|
||||
for (i, &neighbor_id) in neighbors.iter().enumerate() {
|
||||
let attention_weight = if sum_exp > 0.0 {
|
||||
exp_scores[i] / sum_exp
|
||||
} else {
|
||||
1.0 / neighbors.len() as f32
|
||||
};
|
||||
|
||||
node_attentions.push((neighbor_id.index(), attention_weight));
|
||||
|
||||
// Store attention weights for inspection
|
||||
self.attention_weights
|
||||
.insert((node_id.index(), neighbor_id.index()), attention_weight);
|
||||
}
|
||||
}
|
||||
|
||||
attention_map.insert(node_id.index(), node_attentions);
|
||||
}
|
||||
|
||||
Ok(attention_map)
|
||||
}
|
||||
}
|
||||
|
||||
impl GNNLayer for GATLayer {
|
||||
fn forward(&mut self, graph: &Graph) -> Result<Vec<Tensor>> {
|
||||
if graph.node_count() == 0 {
|
||||
return Err(GeomError::ForwardPassFailed {
|
||||
reason: "Cannot process empty graph".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Compute attention weights
|
||||
let attention_map = self.compute_attention(graph)?;
|
||||
|
||||
let mut outputs = Vec::with_capacity(graph.node_count());
|
||||
|
||||
for node_id in graph.nodes() {
|
||||
let node_idx = node_id.index();
|
||||
|
||||
if let Some(attentions) = attention_map.get(&node_idx) {
|
||||
if attentions.is_empty() {
|
||||
// Isolated node - transform its own features
|
||||
let node_features = graph.node(node_id).unwrap().features();
|
||||
let node_2d = node_features.reshape([1, self.input_dim]).map_err(|e| {
|
||||
GeomError::ForwardPassFailed {
|
||||
reason: format!("Self feature reshape failed: {e}"),
|
||||
}
|
||||
})?;
|
||||
|
||||
let transformed =
|
||||
node_2d
|
||||
.matmul(&self.weight)
|
||||
.map_err(|e| GeomError::ForwardPassFailed {
|
||||
reason: format!("Self transformation failed: {e}"),
|
||||
})?;
|
||||
let biased =
|
||||
transformed
|
||||
.add(&self.bias)
|
||||
.map_err(|e| GeomError::ForwardPassFailed {
|
||||
reason: format!("Bias addition failed: {e}"),
|
||||
})?;
|
||||
|
||||
// Reshape back to [output_dim] from [1, output_dim]
|
||||
let output = biased.reshape([self.output_dim]).map_err(|e| {
|
||||
GeomError::ForwardPassFailed {
|
||||
reason: format!("Output reshape failed: {e}"),
|
||||
}
|
||||
})?;
|
||||
|
||||
outputs.push(output);
|
||||
} else {
|
||||
// Aggregate neighbor features with attention weights
|
||||
let mut aggregated = None;
|
||||
|
||||
for &(neighbor_idx, attention_weight) in attentions {
|
||||
// Find neighbor node by index
|
||||
let neighbor_id = graph
|
||||
.nodes()
|
||||
.into_iter()
|
||||
.find(|id| id.index() == neighbor_idx)
|
||||
.ok_or_else(|| GeomError::ForwardPassFailed {
|
||||
reason: format!("Neighbor node {neighbor_idx} not found"),
|
||||
})?;
|
||||
|
||||
let neighbor_features = graph.node(neighbor_id).unwrap().features();
|
||||
let neighbor_2d =
|
||||
neighbor_features
|
||||
.reshape([1, self.input_dim])
|
||||
.map_err(|e| GeomError::ForwardPassFailed {
|
||||
reason: format!("Neighbor feature reshape failed: {e}"),
|
||||
})?;
|
||||
|
||||
let transformed = neighbor_2d.matmul(&self.weight).map_err(|e| {
|
||||
GeomError::ForwardPassFailed {
|
||||
reason: format!("Neighbor transformation failed: {e}"),
|
||||
}
|
||||
})?;
|
||||
|
||||
let weighted = transformed
|
||||
.mul(
|
||||
&Tensor::from_data(
|
||||
vec![attention_weight],
|
||||
[1],
|
||||
&rtx_tensor::Device::default(),
|
||||
)
|
||||
.unwrap(),
|
||||
)
|
||||
.map_err(|e| GeomError::ForwardPassFailed {
|
||||
reason: format!("Attention weighting failed: {e}"),
|
||||
})?;
|
||||
|
||||
aggregated = match aggregated {
|
||||
None => Some(weighted),
|
||||
Some(acc) => Some(acc.add(&weighted).map_err(|e| {
|
||||
GeomError::ForwardPassFailed {
|
||||
reason: format!("Aggregation failed: {e}"),
|
||||
}
|
||||
})?),
|
||||
};
|
||||
}
|
||||
|
||||
let output = if let Some(agg) = aggregated {
|
||||
agg.add(&self.bias)
|
||||
.map_err(|e| GeomError::ForwardPassFailed {
|
||||
reason: format!("Bias addition failed: {e}"),
|
||||
})?
|
||||
} else {
|
||||
return Err(GeomError::ForwardPassFailed {
|
||||
reason: "No aggregated features computed".to_string(),
|
||||
});
|
||||
};
|
||||
|
||||
// Apply ELU activation (approximated with ReLU for simplicity)
|
||||
let activated = output.relu().map_err(|e| GeomError::ForwardPassFailed {
|
||||
reason: format!("Activation failed: {e}"),
|
||||
})?;
|
||||
|
||||
// Reshape back to [output_dim] from [1, output_dim]
|
||||
let output_1d = activated.reshape([self.output_dim]).map_err(|e| {
|
||||
GeomError::ForwardPassFailed {
|
||||
reason: format!("Output reshape failed: {e}"),
|
||||
}
|
||||
})?;
|
||||
|
||||
outputs.push(output_1d);
|
||||
}
|
||||
} else {
|
||||
return Err(GeomError::ForwardPassFailed {
|
||||
reason: format!("No attention computed for node {node_idx}"),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(outputs)
|
||||
}
|
||||
|
||||
fn parameters(&self) -> Vec<Tensor> {
|
||||
vec![
|
||||
self.weight.clone(),
|
||||
self.attention_weight.clone(),
|
||||
self.bias.clone(),
|
||||
]
|
||||
}
|
||||
|
||||
fn enable_gradients(&mut self, enabled: bool) {
|
||||
self.gradients_enabled = enabled;
|
||||
}
|
||||
|
||||
fn has_computed_gradients(&self) -> bool {
|
||||
self.computed_gradients
|
||||
}
|
||||
|
||||
fn backward(&mut self, grad_outputs: &[Tensor]) -> Result<Vec<Tensor>> {
|
||||
if !self.gradients_enabled {
|
||||
return Err(GeomError::BackwardPassFailed {
|
||||
reason: "Gradients are not enabled for this layer".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
if grad_outputs.is_empty() {
|
||||
return Err(GeomError::BackwardPassFailed {
|
||||
reason: "No gradient outputs provided".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Simplified gradient computation
|
||||
let weight_grad =
|
||||
Tensor::zeros(self.weight.shape().dims(), &rtx_tensor::Device::default()).unwrap();
|
||||
let attention_grad = Tensor::zeros(
|
||||
[self.attention_weight.shape().dims()[0]],
|
||||
&rtx_tensor::Device::default(),
|
||||
)
|
||||
.unwrap();
|
||||
let bias_grad = Tensor::zeros([self.output_dim], &rtx_tensor::Device::default()).unwrap();
|
||||
|
||||
self.computed_gradients = true;
|
||||
|
||||
Ok(vec![weight_grad, attention_grad, bias_grad])
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user