Skip to content

Model

Bases: BaseModel, DataMixin

Model provide training functionality with model.fit(...) inspired from Keras

Examples:

model = Model(cnn)
model.compile("crossentropyloss", "adam", learning_rate=1e-3, metrics="accuracy")
model.fit(autodataset)

Parameters:

Name Type Description Default
learner Union[nn.Module, Any]

Trainable model

required
device Optional[str]

auto | cpu | gpu | mps

None
precision Any

Numerical precision value, could be 32 | 16 | "b16"

32
strategy Optional[str]

Strategy for distributed training (ddp | ddp_spawn | deepspeed | fsdp)

None
accelerator_config dict

Accelerator config

None

compile(loss='crossentropyloss', optimizer='adam', learning_rate=0.0003, metrics=None, loss_config=None, optimizer_config=None)

Compile loss function, optimizer and metrics Example:

model = Model(net)
model.compile(loss="crossentropyloss", optimizer="adam", learning_rate=1e-3, metrics="accuracy")
You can also compile optimizer by passing class as argument-
model.compile(optimizer=torch.optim.SGD, learning_rate=1e-3, optimizer_configs = {"momentum":0.9})
To see a list of available losses and metrics-
from gradsflow.models import available_losses, available_metrics
print(available_losses())
print(available_metrics())

Parameters:

Name Type Description Default
loss Union[str, nn.Module]

name of loss, torch Loss class object or any functional method. See available_losses()

'crossentropyloss'
optimizer Union[str, Callable]

optimizer name or torch.optim.Optimizer Class

'adam'
learning_rate float

defaults to 1e-3

0.0003
metrics METRICS_TYPE

list of metrics to calculate. See available_metrics()

None
loss_config Optional[dict]

Dict config if any to pass to loss function

None
optimizer_config Optional[dict]

Dict config if any to pass to Optimizer

None

epoch()

Train & Validate Model for specified number of epochs with on_epoch_* callback method.

fit(autodataset, max_epochs=1, steps_per_epoch=None, callbacks=None, resume=True, show_progress=True, progress_kwargs=None)

Analogous to Keras model.fit(...) API, it trains the model for specified epochs and returns Tracker object

Examples:

autodataset = AutoDataset(train_dataloader, val_dataloader)
model = Model(cnn)
model.compile("crossentropyloss", "adam", learning_rate=1e-3)
model.fit(autodataset)

Parameters:

Name Type Description Default
autodataset AutoDataset

AutoDataset object encapsulate dataloader and datamodule

required
max_epochs int

number of epochs to train

1
steps_per_epoch Optional[int]

Number of steps trained in a single current_epoch

None
callbacks Optional[Union[List[Callback], Callback, str, List[str]]]

Callback object or string

None
resume bool

Resume training from the last current_epoch

True
show_progress bool

Enable to show training progress

True
progress_kwargs

Arguments for rich.progress

None

Returns:

Type Description
Tracker

Tracker object

train_one_epoch(train_dataloader)

Train model for a single epoch with train_dataloader. CallbackRunner will call on_train_step_* method for each training step. Training will break if step >= steps_per_epoch.

val_one_epoch(val_dataloader)

Validate model for a single epoch with val_dataloader. CallbackRunner will call on_val_step_* method for each validation step.


Last update: October 3, 2021