Skorch

Integrate W&B with Skorch to log scikit-learn compatible PyTorch model training metrics and hyperparameters.

This page shows you how to use W&B with Skorch so you can track Skorch model training without writing custom logging code. When you integrate the two, W&B automatically logs the model with the best performance, along with all model performance metrics, the model topology, and compute resources after each epoch. W&B automatically logs every file you save in wandb_run.dir.

For more information, see this example run.

Parameters

The following table lists the parameters that the WandbLogger callback accepts.

Parameter	Type	Description
wandb_run	wandb.wandb_run. Run	The W&B run used to log data.
save_model	bool (default=True)	Whether to save a checkpoint of the best model and upload it to your run on W&B.
keys_ignored	str or list of str (default=None)	Key or list of keys not to log to TensorBoard. In addition to the keys you provide, W&B ignores keys such as those starting with event_ or ending with _best by default.

Example code

The following examples show end-to-end usage of WandbLogger with Skorch:

• Colab: A simple demo to try the integration.
• Step-by-step guide: A walkthrough for tracking your Skorch model performance.

# Install wandb
pip install wandb

import wandb
from skorch.callbacks import WandbLogger

# Create a wandb run
wandb_run = wandb.init()

# Log hyper-parameters (optional)
wandb_run.config.update({"learning rate": 1e-3, "batch size": 32})

net = NeuralNet(..., callbacks=[WandbLogger(wandb_run)])
net.fit(X, y)

Method reference

The following table lists the callback methods that WandbLogger provides and when Skorch invokes each one.

Method	Description
initialize()	(Re-)Set the initial state of the callback.
on_batch_begin(net[, X, y, training])	Called at the beginning of each batch.
on_batch_end(net[, X, y, training])	Called at the end of each batch.
on_epoch_begin(net[, dataset_train, …])	Called at the beginning of each epoch.
on_epoch_end(net, **kwargs)	Log values from the last history step and save the best model.
on_grad_computed(net, named_parameters[, X, …])	Called once per batch after gradients are computed but before an update step is performed.
on_train_begin(net, **kwargs)	Log model topology and add a hook for gradients.
on_train_end(net[, X, y])	Called at the end of training.