xtrain

XTRAIN: a tiny library for training Flax models.

Design goals:

Help avoiding boiler-plate code
Minimal functionality and dependency
Agnostic to hardware configuration (e.g. GPU->TPU)

General workflow

Step 1: define your model

class MyFlaxModule(nn.Module):
  @nn.compact
  def __call__(self, x):
    ...

Step 2: define loss function

def my_loss_func(batch, prediction):
    x, y_true = batch
    loss = ....
    return loss

Step 3: create an iterator that supplies training data

my_data = zip(sequence_of_inputs, sequence_of_labels)

Step 4: train

# create and initialize a Trainer object
trainer = xtrain.Trainer(
  model = MyFlaxModule(),
  losses = my_loss_func,
  optimizer = optax.adam(1e-4),
)

train_iter = trainer.train(my_data) # returns a iterable object

# iterate the train_iter trains the model
for epoch in range(3):
  for model_out in train_iter:
    pass
  print(train_iter.loss_logs)
  train_iter.reset_loss_logs()

Training data format

tensowflow Dataset
torch dataloader
generator function
other python iterable that produce numpy data

Checkpointing

train_iter is orbax compatible.

import orbax.checkpoint as ocp
ocp.StandardCheckpointer().save(cp_path, args=ocp.args.StandardSave(train_iter))

Freeze submodule

train_iter.freeze("submodule/Dense_0/kernel")

Simple batch parallelism on multiple device

# Add a new batch dim to you dataset
ds = ds.batch(8)
# create trainer with the Distributed strategy
trainer_iter = xtrain.Trainer(model, losses, optimizer, strategy=xtrain.Distributed).train(ds)

API documentation

https://jiyuuchc.github.io/xtrain/

Metadata

Release files for xtrain 0.4.3

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for xtrain 0.4.3
File	Size	Uploaded
xtrain-0.4.3.tar.gz	12.8 kB	Dec 18, 2024	Details

Built distribution (wheel)

Table of built distributions (wheels) for xtrain 0.4.3
File	Interpreter	ABI	Platform	Reset
xtrain-0.4.3-py3-none-any.whl 14.4 kB Dec 18, 2024	Python 3	none	any	Details

Total release size: 27.2 kB

Release files / xtrain-0.4.3.tar.gz

Download URL	xtrain-0.4.3.tar.gz
Size	12.8 kB
Tags	Source
SHA-256 checksum How to use checksums	`3793ae065a67dc85b54e557b121c2b323093b278bc120e150a0426440b654d31`
BLAKE2b-256 checksum How to use checksums	`4e1c5e866f402e07463a2acbbecfa2c4f9e038ec2b0f39b83f2bd705d8918da2`
Upload date	Dec 18, 2024
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`poetry/1.8.5 CPython/3.10.15 Linux/6.8.0-1017-azure`

Release files / xtrain-0.4.3-py3-none-any.whl

Download URL	xtrain-0.4.3-py3-none-any.whl
Size	14.4 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`5dba9a035af9c7355246305b2ecf5b43602768ae4e8553b983cc5c22cef928a3`
BLAKE2b-256 checksum How to use checksums	`486ae9e407f49be9b92d65429c81042a8eb5b3333ca4c5482c797ec37581ad24`
Upload date	Dec 18, 2024
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`poetry/1.8.5 CPython/3.10.15 Linux/6.8.0-1017-azure`