Installation

All you have to do for pystematic to find the plugin is to install it:

$ pip install pystematic-torch

Example

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)

Documentation

Reference documentation is available at https://pystematic-torch.readthedocs.io.