【Rust日报】2024-03-30 krnl v0.1.0: 安全, 可移植, 高性能计算内核

#[derive(Layer, Forward)]
#[autograph(forward(Variable4, Output=Variable2))]
struct LeNet5 {
    conv1: Conv2,
    relu1: Relu,
    pool1: MaxPool2,
    conv2: Conv2,
    relu2: Relu,
    pool2: MaxPool2,
    flatten: Flatten,
    dense1: Dense,
    relu3: Relu,
    dense2: Dense,
    relu4: Relu,
    dense3: Dense,
}
impl LeNet5 {

fn new(device: Device, scalar_type: ScalarType) -> Result<Self> {

let conv1 = Conv2::builder()

.device(device.clone())

.scalar_type(scalar_type)

.inputs(1)

.outputs(6)

.filter([5, 5])

.build()?;

let relu1 = Relu;

let pool1 = MaxPool2::builder().filter([2, 2]).build();

let conv2 = Conv2::builder()

.device(device.clone())

.scalar_type(scalar_type)

.inputs(6)

.outputs(16)

.filter([5, 5])

.build()?;

let relu2 = Relu;

let pool2 = MaxPool2::builder().filter([2, 2]).build();

let flatten = Flatten;

let dense1 = Dense::builder()

.device(device.clone())

.scalar_type(scalar_type)

.inputs(16 * 4 * 4)

.outputs(128)

.build()?;

let relu3 = Relu;

let dense2 = Dense::builder()

.device(device.clone())

.scalar_type(scalar_type)

.inputs(128)

.outputs(84)

.build()?;

let relu4 = Relu;

let dense3 = Dense::builder()

.device(device.clone())

.scalar_type(scalar_type)

.inputs(84)

.outputs(10)

.bias(true)

.build()?;

Ok(Self {

conv1,

relu1,

pool1,

conv2,

relu2,

pool2,

flatten,

dense1,

relu3,

dense2,

relu4,

dense3,

})

}

}
let mut model = LeNet5::new(device.clone(), ScalarType::F32)?;

model.init_parameter_grads()?;

let y = model.forward(x)?;

let loss = y.cross_entropy_loss(t)?;

loss.backward()?;

model.update(learning_rate, &optimizer)?;