Symbolic Neural Modeling¶
This notebook illustrates how to use PyGX to symbolically manipulate PyTorch layers for neural modeling.
!pip install pygx torch
Symbolizing PyTorch Layers¶
Before we can manipulate the combination of PyTorch layers, we symbolize the
torch.nn layer classes via pg.symbolize. This keeps each class's original
behavior -- the symbolized layers are still real nn.Modules you can run --
and adds a symbolic representation on top.
import torch
from torch import nn
import pygx as pg
# Symbolize PyTorch layers.
Sequential = pg.symbolize(nn.Sequential)
Conv2d = pg.symbolize(nn.Conv2d)
Linear = pg.symbolize(nn.Linear)
Flatten = pg.symbolize(nn.Flatten)
ReLU = pg.symbolize(nn.ReLU)
Creating a Symbolic Model¶
By using the symbolic layer classes, we can create a symbolic neural model for 2D image classification with 10 classes, sized for 28x28 single-channel images.
nn.Conv2d and nn.Linear take their input width explicitly, so each layer
states both what it consumes and what it produces.
def create_model():
return Sequential(
Conv2d(1, 16, (5, 5)),
ReLU(),
Conv2d(16, 32, (3, 3)),
ReLU(),
Flatten(),
Linear(32 * 22 * 22, 10),
)
model = create_model()
# The symbolized PyTorch layers can be printed in human readable form.
# For clarity, we hide the default values of the layers.
print(model.format(exclude_defaults=True))
Sequential(
args = [
0 : Conv2d(
in_channels = 1,
out_channels = 16,
kernel_size = (5, 5)
),
1 : ReLU(),
2 : Conv2d(
in_channels = 16,
out_channels = 32,
kernel_size = (3, 3)
),
3 : ReLU(),
4 : Flatten(),
5 : Linear(
in_features = 15488,
out_features = 10
)
]
)
Manipulating Models¶
model is a regular nn.Module -- it runs like any other PyTorch model:
images = torch.randn(2, 1, 28, 28)
print(model(images).shape)
torch.Size([2, 10])
What if we want to upscale the model by increasing the number of output channels by 2?
def double_width(k, v, p):
"""A rebind rule for doubling the output channels of Conv2d layers.
Args:
k: A `pg.KeyPath` object representing the location of current node.
v: The value of current node.
p: The parent of current node.
Returns:
The output value for current node.
"""
if isinstance(p, Conv2d) and k.key == 'out_channels':
return 2 * v
return v
# Rebind allows the users to manipulate a symbolic object by
# rules. It mutates the object in place and returns it.
wider = create_model().sym_rebind(double_width)
print(wider.format(exclude_defaults=True))
Sequential(
args = [
0 : Conv2d(
in_channels = 1,
out_channels = 32,
kernel_size = (5, 5)
),
1 : ReLU(),
2 : Conv2d(
in_channels = 16,
out_channels = 64,
kernel_size = (3, 3)
),
3 : ReLU(),
4 : Flatten(),
5 : Linear(
in_features = 15488,
out_features = 10
)
]
)
Widening only the outputs leaves each layer's declared input width behind, so
the second Conv2d above still says in_channels = 16 while its predecessor
now emits 32. A rule that rewrites both sides keeps the model runnable:
def double_width_consistently(k, v, p):
# The first convolution reads the image itself, so its input width is
# fixed by the data rather than by the previous layer.
if isinstance(p, Conv2d) and k.key == 'out_channels':
return 2 * v
if isinstance(p, Conv2d) and k.key == 'in_channels' and v != 1:
return 2 * v
# The classification head consumes the flattened final feature map.
if isinstance(p, Linear) and k.key == 'in_features':
return 2 * v
return v
wider = create_model().sym_rebind(double_width_consistently)
print(wider(images).shape)
torch.Size([2, 10])
What if we want to remove the ReLU activations?
Deleting an element from a list is not something a rebind rule can express --
sym_rebind visits values in place, and a list has no "absent" slot to leave
behind. Instead we assign the whole list of layers at once. pg.maybe_clone
detaches each kept layer from the model it currently belongs to, since a
symbolic node can only be held in one position at a time.
model = create_model()
kept = [
pg.maybe_clone(layer)
for layer in model.sym_getattr('args')
if not isinstance(layer, ReLU)
]
print(model.sym_rebind(args=kept).format(exclude_defaults=True))
Sequential(
args = [
0 : Conv2d(
in_channels = 1,
out_channels = 16,
kernel_size = (5, 5)
),
1 : Conv2d(
in_channels = 16,
out_channels = 32,
kernel_size = (3, 3)
),
2 : Flatten(),
3 : Linear(
in_features = 15488,
out_features = 10
)
]
)
What if we want to change the number of classes for the classification head from 10 to 100?
model = create_model()
# Query the last Linear layer in the model, and modify its out_features to 100.
result = pg.query(model, where=lambda v: isinstance(v, Linear))
classification_head_location = list(result.keys())[-1]
model.sym_rebind({
f'{classification_head_location}.out_features': 100
})
print(model.format(exclude_defaults=True))
Sequential(
args = [
0 : Conv2d(
in_channels = 1,
out_channels = 16,
kernel_size = (5, 5)
),
1 : ReLU(),
2 : Conv2d(
in_channels = 16,
out_channels = 32,
kernel_size = (3, 3)
),
3 : ReLU(),
4 : Flatten(),
5 : Linear(
in_features = 15488,
out_features = 100
)
]
)
Because the symbolized layers stay real nn.Modules, editing the symbolic
representation rebuilds the underlying module: the model now emits 100 logits
instead of 10, and the head's weight matrix has been reallocated to match.
print(model(images).shape)
print(model[5].weight.shape)
torch.Size([2, 100]) torch.Size([100, 15488])