rust-candle学习笔记11-实现一个简单的自注意力

参考:about-pytorch

定义ScaledDotProductAttention结构体:

rust 复制代码
use candle_core::{Result, Device, Tensor};
use candle_nn::{Linear, Module, linear_no_bias, VarMap, VarBuilder, ops};

struct ScaledDotProductAttention {
    wq: Linear,
    wk: Linear,
    wv: Linear,
    d_model: Tensor,
    device: Device,
}

为ScaledDotProductAttention结构体实现new方法:

rust 复制代码
impl ScaledDotProductAttention {
    fn new(vb: VarBuilder, embedding_dim: usize, out_dim: usize, device: Device) -> Result<Self> {
        Ok(Self { 
            wq: linear_no_bias(embedding_dim, out_dim, vb.pp("wq"))?, 
            wk: linear_no_bias(embedding_dim, out_dim, vb.pp("wk"))?, 
            wv: linear_no_bias(embedding_dim, out_dim, vb.pp("wv"))?,
            d_model: Tensor::new(embedding_dim as f32, &device)?,
            device,
        })
    }
}

为结构体实现Module的forward trait:

rust 复制代码
impl Module for ScaledDotProductAttention {
    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        let q = self.wq.forward(xs)?;
        let k = self.wk.forward(xs)?;
        let v = self.wv.forward(xs)?;
        let attn_score = q.matmul(&k.t()?)?;
        let attn_score = attn_score.broadcast_div(&self.d_model.sqrt()?)?;
        let dim = attn_score.rank() - 1;
        let attn_weights = ops::softmax(&attn_score, dim)?;
        let attn_output = attn_weights.matmul(&v)?;
        Ok(attn_output)
    }
}

融合qkv实现:

定义ScaledDotProductAttentionFusedQKV结构体:

rust 复制代码
struct ScaledDotProductAttentionFusedQKV {
    w_qkv: Linear,
    d_model: Tensor,
    device: Device,
}

为结构体实现new方法:

rust 复制代码
impl ScaledDotProductAttentionFusedQKV {
    fn new(vb: VarBuilder, embedding_dim: usize, out_dim: usize, device: Device) -> Result<Self> {
        Ok(Self { 
            w_qkv: linear_no_bias(embedding_dim, 3*out_dim, vb.pp("w_qkv"))?,
            d_model: Tensor::new(embedding_dim as f32, &device)?,
            device,
        })
    }
}

为结构体实现forward trait:

rust 复制代码
impl Module for ScaledDotProductAttentionFusedQKV {
    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        let qkv = self.w_qkv.forward(xs)?;
        let (batch_size, seq_len, _) = qkv.dims3()?;
        let qkv = qkv.reshape((batch_size, seq_len, 3, ()))?;
        let q = qkv.get_on_dim(2, 0)?;
        let q = q.reshape((batch_size, seq_len, ()))?;
        let k = qkv.get_on_dim(2, 1)?;
        let k = k.reshape((batch_size, seq_len, ()))?;
        let v = qkv.get_on_dim(2, 2)?;
        let v = v.reshape((batch_size, seq_len, ()))?;
        let attn_score = q.matmul(&k.t()?)?;
        let attn_score = attn_score.broadcast_div(&self.d_model.sqrt()?)?;
        let dim = attn_score.rank() - 1;
        let attn_weights = ops::softmax(&attn_score, dim)?;
        let attn_output = attn_weights.matmul(&v)?;
        Ok(attn_output)
    }
}

测试:

rust 复制代码
fn main() -> Result<()> {
    let device = Device::cuda_if_available(0)?;
    let varmap = VarMap::new();
    let vb = VarBuilder::from_varmap(&varmap, candle_core::DType::F32, &device);
    
    let input = Tensor::from_vec(vec![0.43f32, 0.15, 0.89, 
                                                    0.55, 0.87, 0.66,
                                                    0.57, 0.85, 0.64,
                                                    0.22, 0.58, 0.33,
                                                    0.77, 0.25, 0.10,
                                                    0.05, 0.80, 0.55, 
                                                    0.43, 0.15, 0.89, 
                                                    0.55, 0.87, 0.66,
                                                    0.57, 0.85, 0.64,
                                                    0.22, 0.58, 0.33,
                                                    0.77, 0.25, 0.10,
                                                    0.05, 0.80, 0.55], (2, 6, 3), &device)?;
    // let model = ScaledDotProductAttention::new(vb.clone(), 3, 2, device.clone())?;
    let model = ScaledDotProductAttentionFusedQKV::new(vb.clone(), 3, 2, device.clone())?;
    let output = model.forward(&input)?;
    println!("output: {:?}\n", output);
    println!("output: {:?}\n", output.to_vec3::<f32>()?);
    Ok(())
}
相关推荐
使二颗心免于哀伤33 分钟前
《设计模式之禅》笔记摘录 - 10.装饰模式
笔记·设计模式
悠哉悠哉愿意1 小时前
【电赛学习笔记】MaxiCAM 项目实践——与单片机的串口通信
笔记·python·单片机·嵌入式硬件·学习·视觉检测
快乐肚皮1 小时前
ZooKeeper学习专栏(五):Java客户端开发(原生API)详解
学习·zookeeper·java-zookeeper
慕y2741 小时前
Java学习第七十二部分——Zookeeper
java·学习·java-zookeeper
想变成树袋熊1 小时前
【自用】NLP算法面经(6)
人工智能·算法·自然语言处理
岩中竹2 小时前
广东省省考备考——常识:科技常识(持续更新)
笔记
★YUI★2 小时前
学习游戏制作记录(剑投掷技能)7.26
学习·游戏·unity·c#
Olrookie2 小时前
若依前后端分离版学习笔记(三)——表结构介绍
笔记·后端·mysql
蓝桉8022 小时前
opencv学习(图像金字塔)
人工智能·opencv·学习
rannn_1113 小时前
Java学习|黑马笔记|Day23】网络编程、反射、动态代理
java·笔记·后端·学习