peal.adaptors.clarc.split_model¶

peal.adaptors.clarc.split_model(model, split_at, device)[source]¶

Split a model into Sequential feature extractor and downstream head.

Parameters:

model (torch.nn.Module) – Model whose (recursively flattened) children are split.
split_at (int) – Index into the flattened child list; children before it form the feature extractor, the rest the head.
device – Device both parts are moved to.

Returns:

(feature_extractor, downstream_head). For torchvision ResNet a Flatten is inserted before the final fc layer since the flattening is not a child module.

Return type:

(torch.nn.Module, torch.nn.Module)