Train (neural_nets_lib.Ocannl.Train)

module Ops = Ir.Ops

module Tn = Ir.Tnode

module Nd = Ir.Ndarray

module Asgns = Ir.Assignments

module Idx = Ir.Indexing

module Task = Ir.Task

val _get_local_debug_runtime : unit -> (module Minidebug_runtime.Debug_runtime)

module CDSL : sig ... end

module IDX : sig ... end

val run : Context.t -> Context.routine -> Base.unit

val set_materialized : Tn.t -> unit

val forward : 
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Ir.Assignments.comp

Sets the tensor's value as materialized (device-resident, inspectable on demand via the context), and returns the tensor's forward code with a label-derived comment.

val grad_update : 
  ?setup_for_parallel:bool ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Ocannl_tensor.Operation.DSL_modules.Ir.Assignments.comp

Returns the tensor's forward, zeroing gradients, and backprop code wrapped with label-derived comments. Sets the tensor's value as materialized. If setup_for_parallel is true (false by default), sets the parameters and their gradients as "non-local" (on-device).

val sgd_one : 
  learning_rate:Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  ?momentum:Base.Float.t ->
  ?weight_decay:Base.float ->
  ?nesterov:bool ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Ocannl_tensor.Operation.DSL_modules.Ir.Assignments.comp

See: https://github.com/tinygrad/tinygrad/blob/master/tinygrad/nn/optim.py

val sgd_update : 
  learning_rate:Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  ?momentum:Base.Float.t ->
  ?weight_decay:Base.float ->
  ?nesterov:bool ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Asgns.comp

val sequential_loop : 
  f:(unit -> Base.unit) ->
  (Idx.static_symbol * int Base.ref) list ->
  Base.unit

All and only bindings with associated ranges are iterated, with the binding's initial value lost. Bindings without ranges remain at their initial values.

val set_virtual : Tn.t -> unit

val every_non_literal_materialized : 
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Base.unit

Materializes every non-literal embedded tensor node of t (so its value is inspectable on demand via the context). Replaces the old every_non_literal_on_host now that there is no hosted memory mode (gh-ocannl-333).

module Lazy = Utils.Lazy

val to_routine : 
  Context.t ->
  ?output_cd_file:bool ->
  (Base.unit -> Base.unit) Idx.bindings ->
  Asgns.comp ->
  Context.routine

val init_params : 
  ?reinit_all:bool ->
  Context.t ->
  Ir.Indexing.unit_bindings ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Context.t

init_params initializes the parameters of t, via running their forward code or copying from the host as appropriate. If reinit_all is true, all parameters are reinitialized, otherwise only the parameters that are not in ctx.ctx_buffers are initialized.

type example_train_result = {

inputs : Ocannl_tensor.Operation.DSL_modules.Tensor.t;
outputs : Ocannl_tensor.Operation.DSL_modules.Tensor.t;
model_result : Ocannl_tensor.Operation.DSL_modules.Tensor.t;
(*
Do not use model_result for deriving gradients.
*)
infer_callback : Base.float Base.array -> Base.float Base.array;
(*
Computes the output for the given input via the model_result tensor. Note: infer_callback is inefficient as it is not batched.
*)
rev_batch_losses : Base.float Base.list;
rev_epoch_losses : Base.float Base.list;
learning_rates : Base.float Base.list;
used_memory : Base.int;

}

val run_once : 
  ?output_cd_file:bool ->
  ?skip_init:Base.bool ->
  ?reinit_all:bool ->
  ?bindings:(Base.unit -> Base.unit) Idx.bindings ->
  f:(Ocannl_tensor.Operation.DSL_modules.Tensor.t -> Asgns.comp) ->
  Context.t ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Context.t

run_once is a wrapper around init_params that additionally runs code of f t and returns the context. If skip_init is true (false by default), no initialization is performmed. If reinit_all is true (false by default), all parameters are reinitialized, otherwise only the parameters that are not in ctx.ctx_buffers are initialized.

If output_cd_file is true, the global setting output_debug_files_in_build_directory must be true, and the update code is output to a file before shape inference potentially crashes at init_params.

Context-based versions of training functions for the new simplified API

val forward_once : 
  ?output_cd_file:bool ->
  ?skip_init:Base.bool ->
  ?reinit_all:bool ->
  ?bindings:(Base.unit -> Base.unit) Idx.bindings ->
  Context.t ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Context.t

forward_once is a wrapper around run_once that runs the forward code of t.

val update_once : 
  ?output_cd_file:bool ->
  ?skip_init:Base.bool ->
  ?reinit_all:bool ->
  ?bindings:(Base.unit -> Base.unit) Idx.bindings ->
  Context.t ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Context.t

update_once is a wrapper around run_once that runs the gradient update code of t: both forward and backprop.

val ensure_printable : 
  Context.t ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Context.t

val printf : 
  ?here:Ppx_here_lib.position ->
  ?with_grad:Base.bool ->
  ?with_code:Base.bool ->
  ?with_low_level:Base.bool ->
  ?style:Ocannl_tensor.Operation.DSL_modules.Tensor.array_print_style ->
  Context.t ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Base.unit

printf is a wrapper around Tensor.print that assumes ~force:true, and by default sets ~with_code:false, ~with_grad:true, and ~style:`Default. It takes an explicit context and retrieves values on demand (gh-ocannl-333). If the tensor's value is not already materialized in ctx, it is recomputed via the for_print copy trick so real values are still shown.

val printf_tree : 
  ?here:Ppx_here_lib.position ->
  ?with_value:Base.bool ->
  ?with_grad:Base.bool ->
  ?depth:Base.int ->
  Context.t ->
  Ocannl_tensor.Operation.DSL_modules.Tensor.t ->
  Base.unit

printf_tree is a wrapper around Tensor.print_tree that assumes ~force:true, and by default sets ~with_value:true, ~with_grad:true, and ~depth:9. It takes an explicit context and retrieves values on demand (recomputing via for_print if not already materialized).