Implementation of the same neural network in Scala front end

demo3.scala

// Implementation of the same neural network in Scala front end

package lantern
...

def snippet(a: Rep[String]): Rep[Unit] = {
    ...
    val train_loader = new DataLoader(...)

    // initialize all parameters
    val pars = List(...) 
    def forward(input: TensorR, target: Rep[Int]) = { (dummy: TensorR) =>
        val resL1 = (input.view(784).dot(pars(0)) + pars(1)).relu()
        val resL2 = resL1.dot(pars(2)) + pars(3)
        resL2.logSoftmax().nllLoss(target)
    }

    def train(epoch: Rep[Int]) = {
        for ((batch_idx, (data, target)) <- zip(train_loader.length, train_loader)) {
            val loss = forward(data, target)
            val loss_value = grad_loss(loss) // grad_loss triggers backward prop and returns loss_value
            for (par <- pars) step()
            if (((batch_idx + 1) % log_interval) == 0):
                print_time_and_loss()
        }
    }

    for (idx <- epochs) train(idx)
}