This is an extension to pystematic that adds functionality related to running machine learning experiments in pytorch. Its main contribution is the Context
object and related classes. Which provides an easy way to manage all pytorch related objects.
All you have to do for pystematic to find the plugin is to install it:
$ pip install pystematic-torch
Here's a small example that shows how using the Context
object, SmartDataLoader
and Recorder
simplifies setting up and running a training session in pytorch.
import pystematic
@pystematic.experiment
def context_example(params):
ctx = pystematic.torch.Context()
ctx.epoch = 0
ctx.recorder = pystematic.torch.Recorder()
ctx.model = torch.nn.Sequential(
torch.nn.Linear(2, 1),
torch.nn.Sigmoid()
)
ctx.optimzer = torch.optim.SGD(ctx.model.parameters(), lr=0.01)
# We use the smart dataloader so that batches are moved to
# the correct device
ctx.dataloader = pystematic.torch.SmartDataLoader(
dataset=Dataset(),
batch_size=2
)
ctx.loss_function = torch.nn.BCELoss()
ctx.cuda() # Move everything to cuda
# ctx.ddp() # and maybe distributed data-parallel?
if params["checkpoint"]:
# Load checkpoint
ctx.load_state_dict(pystematic.torch.load_checkpoint(params["checkpoint"]))
# Train one epoch
for input, lbl in ctx.dataloader:
# The smart dataloader makes sure the batch is placed on
# the correct device.
output = ctx.model(input)
loss = ctx.loss_function(output, lbl)
ctx.optimzer.zero_grad()
loss.backward()
ctx.optimzer.step()
ctx.recorder.scalar("train/loss", loss)
ctx.recorder.step()
ctx.epoch += 1
# Save checkpoint
pystematic.torch.save_checkpoint(ctx.state_dict(), id=ctx.epoch)
Reference documentation is available at https://pystematic-torch.readthedocs.io.