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_dict
s 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.