纯 Mojo 中的张量和动态神经网络Infermo
开源地址:https://github.com/TilliFe/Infermo
作者描述:
我从头开始重建整个引擎。它现在支持通过内存池、梯度检查点和内存共享进行内存管理。这些优化对于游戏后期进行内核融合至关重要。此外,它现在有一个更好的编程界面。
注意:我暂时减少了可用操作员的数量,并且仓库中可供查看的内容通常较少,但我将尝试在接下来的几天内提供这些内容。此外,我们还需要做一些控制流程,稍后会详细介绍。

Infermo 是一个 Mojo 库,它提供了两个高级功能:
Mojo 目前仅在 CPU 上运行。GPU 支持即将推出!
Infermo 仍然是一个概念验证,如果您遇到任何错误,请随时创建问题或 PR。😊
一个小小的例子
# lets's build a simple neural network that learns to approximate sin(15x)
fn main() raises:
# init params
let W1 = Tensor(shape(1,64)).randhe().requires_grad()
let W2 = Tensor(shape(64,64)).randhe().requires_grad()
let W3 = Tensor(shape(64,1)).randhe().requires_grad()
let b1 = Tensor(shape(64)).randhe().requires_grad()
let b2 = Tensor(shape(64)).randhe().requires_grad()
let b3 = Tensor(shape(1)).randhe().requires_grad()
# training
var avg_loss = Float32(0.0)
let every = 1000
let num_epochs = 1000000
for epoch in range(1,num_epochs+1):
# set input and true values
let input = Tensor(shape(32,1)).randu(0,1).dynamic()
let true_vals = sin(15.0 * input)
# define model architecture
var x = relu(input @ W1 + b1)
x = relu(x @ W2 + b2)
x = x @ W3 + b3
let loss = mse(x,true_vals).forward()
# print progress
avg_loss += loss[0]
if epoch%every == 0:
print("Epoch:",epoch," Avg Loss: ",avg_loss/every)
avg_loss = 0.0
# compute gradients and optimize
loss.backward()
loss.optimize(0.01,"sgd")
# clear graph
loss.clear()
独特之处
- 内存共享
- 梯度检查点
- 在静态和动态图形执行之间进行选择
即将推出。。。
- 更优化的内存管理
- GPU 支持
- 更多运算符、激活、优化器