def train_read

import bendtorch-load@1.3.0.0/main.bend as BendtorchLoad

import Base
import bendtorch@1.3.0.0/main.bend as T

188# The checkpoint layout is determined by the already initialized optimizer state, without a header:
189# SGD: parameters; momentum: parameters, buffer; Adam/AdamW: parameters, m, v, 1 - b1^t, 1 - b2^t; AdamW with packed
190# BF16 moments: parameters, the moments' tree (main.bend's pt_store), 1 - b1^t, 1 - b2^t, the whole file read.
191def train_read(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer<P, X, Y, S>, +path: String, +base: U32, +copies: U32, +tail: U32, tr: T.Train<P>) -> IO(T.Train<P>):
192  do IO<T.Train<P>>:
193    require(U32.is_le(base, U32.div(U32.sub(4294967295, tail), copies)), "checkpoint exceeds 2^32 - 1 numbers")
194    +n : U32 = ((copies * base : U32) + tail : U32)
195    a : Array<F32> <- read(path, n)
196    return T.Train.restore(~P, ~X, ~Y, ~S, ~m, tr, a)

Uses (18)

Used by (1)

bendtorch-load@1.3.0.0/main:train