SGD

import std.stdio;
import numir;
import grain.autograd; // : Variable, HostStorage;

{
    auto mlp = MLP!(float, HostStorage)(3);
    mlp.fc1.weight.grad[0] = 1.0;
    mlp.zeroGrad();
    assert(mlp.fc1.weight.grad[0] == 0.0);

    auto sgd = SGD!(typeof(mlp))(mlp, 0.5);
    mlp.fc1.weight.data.zero_();
    mlp.fc1.weight.grad = [[1.0f, 0.0f, 0.0f], [0.0f, 0.0f, 0.0f]].variable
        .data;
    sgd.update();
    assert(mlp.fc1.weight.sliced == [[-0.5, 0.0, 0.0], [0.0, 0.0, 0.0]]);
}
version (grain_cuda) {
    auto mlp = MLP!(float, DeviceStorage)(3);
    mlp.fc1.weight.grad = [[1.0f, 0.0f, 0.0f], [0.0f, 0.0f,
        0.0f]].variable.to!DeviceStorage.data;
    mlp.zeroGrad();
    assert(mlp.fc1.weight.to!HostStorage.gradSliced == [[0.0, 0.0, 0.0], [0.0,
            0.0, 0.0]]);

    auto sgd = SGD!(typeof(mlp))(mlp, 0.5);
    mlp.fc1.weight.data.zero_();
    mlp.fc1.weight.grad = [[1.0f, 0.0f, 0.0f], [0.0f, 0.0f,
        0.0f]].variable.to!DeviceStorage.data;
    sgd.update();
    assert(mlp.fc1.weight.to!HostStorage.sliced == [[-0.5, 0.0, 0.0], [0.0, 0.0,
            0.0]]);
}