Training And Checkpoints

rf.Model is a Lightning module. relflow supplies schema-based computation, losses, metrics, and tensorfield callbacks. Lightning controls the training loop, devices, precision, logging, and training checkpoints.

To learn from unlabeled records and then adapt to a labeled task, follow Pretraining And Fine-Tuning.

Fit A Model

Start with a constructed model and independent Arrow train_table and validation_table inputs. The model needs at least one reconstruction objective, such as mask=True on a label, and an optimizer before fitting.

import lightning.pytorch as lit

import relflow as rf

model.optimizer = rf.adamw(learning_rate=1e-3)
datamodule = rf.ArrowDataModule(
    model=model,
    train=train_table,
    validate=validation_table,
)
checkpoint = rf.RollbackCheckpoint(
    dirpath="checkpoints",
    monitor="loss/validate",
    mode="min",
    save_top_k=1,
    save_last=True,
)
trainer = lit.Trainer(max_epochs=50, callbacks=[checkpoint])
trainer.fit(model=model, datamodule=datamodule)

For eager Polars frames, use rf.PolarsDataModule(model=model, train=train_frame, validate=validation_frame). It converts the frames to Arrow at ingress. See Data Modules for source and split choices.

Pass an optimizer instance or a factory receiving the model, either through Model(optimizer=...) or model.optimizer. A factory collects the current parameters when Lightning configures optimization. scheduler accepts a scheduler, a Lightning scheduler configuration, or a factory receiving (model, optimizer).

RollbackCheckpoint restores the best saved model state when fitting ends. It requires full checkpoints and at least one saved checkpoint. Use Lightning’s ordinary ModelCheckpoint if the in-memory model should stay at its final training state. Read Evaluation before choosing a monitor.

Choose The Loop

Call	Required data split	Purpose
trainer.fit(...)	train; optionally validate	Optimize, optionally validating each epoch.
trainer.validate(...)	validate	Evaluate validation objectives.
trainer.test(...)	test	Evaluate held-out objectives.
trainer.predict(...)	predict	Produce predictions through Lightning.

Stages also control masking and field-owned learning, such as vocabulary and normalizer updates. Use the corresponding split rather than repurposing a training loader for inference. For bounded requests, the separate model.predict(observations) convenience returns an Arrow table without a Trainer; see Prediction Output.

Set accelerator, devices, precision, and distributed strategy on the Trainer. relflow installs its own required callbacks automatically. Compile selected model regions before the loop with model.compile() if desired; Performance describes exactly what this compiles.

Save And Load

Use a relflow artifact to retain the schema and learned model state:

model.save("artifacts/customer.rf")
loaded = rf.Model.load("artifacts/customer.rf")
predictions = loaded.predict(observations)

The artifact contains the state dictionary, schema, batch size, and producing relflow version. State includes field-owned buffers such as vocabularies and normalizers. load constructs the model on CPU. Move it to the intended device for direct inference; Lightning owns placement when using its loops.

A Model.save artifact does not contain optimizer or scheduler state, Trainer progress, data sources, or processor configuration. Supply those explicitly when fitting a loaded model. Compilation is transient and must be requested again after loading.

Resume Training

Use a full Lightning checkpoint to resume optimizer/scheduler state and loop progress. Recreate the runtime configuration and pass the checkpoint to fit:

resumed = rf.Model.load("checkpoints/last.ckpt")
resumed.optimizer = rf.adamw(learning_rate=1e-3)
resume_data = rf.ArrowDataModule(
    model=resumed,
    train=train_table,
    validate=validation_table,
)
trainer = lit.Trainer(max_epochs=100)
trainer.fit(
    model=resumed,
    datamodule=resume_data,
    ckpt_path="checkpoints/last.ckpt",
)

Recreate any scheduler and keep optimizer configuration consistent with the saved run. rf.Model.load restores model state; the ckpt_path argument is what tells Lightning to resume training state.