Burn 是一个新的全面的动态深度学习框架,以极高的灵活性、计算效率和可移植性为其主要目标。
性能
我们相信深度学习框架的目标是将计算转化为有用的智能,因此我们将性能视为Burn的核心支柱。
我们努力通过利用下面描述的多种优化技术来实现最高效率。
自动内核融合 💥
使用Burn意味着在任何后端对您的模型进行优化。
在可能的情况下,我们提供一 种自动和动态创建自定义内核的方式, 最小化在不同内存空间之间的数据重定位,当内存迁移是瓶颈时非常有用。
例如,您可以使用高级张量 API(请参见下面的Rust 代码片段)编写自己的GELU 激活函数。
fn gelu_custom < B : Backend , const D : usize > ( x : Tensor < B , D > ) -> Tensor < B , D > {
let x = x . clone ( ) * ( ( x / SQRT_2 ) . erf ( ) + 1 ) ;
x / 2
}
然后,在运行时, 将自动为您的特定实现创建一个自定义低级内核,并将与手工编写的 GPU 实现相媲美。
内核由大约 60 行 WGSL WebGPU 着色语言组成,这是一种非常冗长的低级着色语言,您可能不希望用其来编写深度学习模型!
就目前而言,我们的融合策略仅为我们自己的 WGPU 后端实现,并且仅支持部分操作。
我们计划很快添加更多操作,并将这种技术扩展到未来的其他内部后端。
异步执行 ❤️🔥
对由 Burn 团队从头开始开发的后端,使用了一种异步执行风格,这允许执行各种优化,例如先前提到的自动内核融合。 异步执行还确保框架的正常执行不会阻塞模型计算,这意味着框架的开销不会显著影响执行速度。 反之,模型中的强烈计算不会干扰框架的响应。 有关我们异步后端的更多信息,