Using PyTorch Modules with Transforms

Remember from the section on creating Transforms that PyTorch Modules can be converted into Transforms. PADL has built-in capabilities that make it convenient to work with them.

PADL Automatically Saves PyTorch State-Dicts

When using padl.save() to save a Transform, PADL automatically stores state_dicts of the contained PyTorch layers along with the code. When loading with padl.load(), these are used to initialize the layers’ parameters.

Devices

Use pd_to() to send all PyTorch layers in a Transform to a device:

my_pipeline = (
    preprocess
    >> batch
    >> layer1
    >> layer2
)

my_pipeline.pd_to('cuda:1')

Accessing Layers and Parameters

All PyTorch layers in a Transform (this includes those nested within Pipelines) are accessible via pd_layers():

layers = list(my_pipeline.pd_layers())

Use pd_parameters() to iterate over all parameters in the Transform. This can be used to initialize a PyTorch optimizer and create a training loop:

predict = (
    preprocess
    >> batch
    >> layer1
    >> layer2
)
train_model = predict / batch >> loss

optimizer = torch.optim.Adam(train_model.pd_parameters(), lr=LR)

for loss in model.train_apply(TRAIN_DATA, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS):
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Weight Sharing

It is possible to use the same PyTorch layer Transforms in multiple Pipelines. Those Transforms then share weights. This can be used, for instance, to create a train Pipeline and an inference Pipeline:

layer = MyPyTorchLayer()

train_pipeline = (
    load
    >> preprocess
    >> augment
    >> batch
    >> layer
    >> loss
)

infer_pipeline = (
    load
    >> preprocess
    >> batch
    >> layer
    >> lookup
)

After some training, both train_pipeline and infer_pipeline will have adapted parameters.

Learn about extras and utilities in the next section.