7.4 General Models: nn_module()

nn_module() is “factory function” for building models of arbitrary complexity. More flexible than the sequential model. Use to define:

• weight initialization

• model structure (forward pass), including identification of model parameters using nn_parameter() .

Example:

my_linear <- nn_module(
  initialize = function(in_features, out_features){
    self$w <- nn_parameter(torch_randn(in_features, out_features)) # random normal
    self$b <- nn_parameter(torch_zeros(out_features))              # zeros
  },
  forward = function(input){
    input$mm(self$w) + self$b
  }
)

Next instantiate the model with input and output dimensions:

l <- my_linear(7, 1)
l

## An `nn_module` containing 8 parameters.
## 
## ── Parameters ──────────────────────────────────────────────────────────────────
## • w: Float [1:7, 1:1]
## • b: Float [1:1]

Apply the model to random data (just like we did in the previous section):

output <- l(torch_randn(5, 7))
output

## torch_tensor
## -4.0846
##  0.1639
## -0.3429
## -0.9869
##  1.4769
## [ CPUFloatType{5,1} ][ grad_fn = <AddBackward0> ]

That was the forward pass. Let’s define a (dummy) loss function and compute the gradient:

loss <- output$mean()
loss$backward() # compute gradient
l$w$grad #inspect result

## torch_tensor
##  0.4843
## -0.2944
##  0.0236
## -0.2049
## -0.5343
##  0.1681
##  0.0149
## [ CPUFloatType{7,1} ]