Package: torchextractor
torchextractor: PyTorch Intermediate Feature ExtractionIntroductionToo many times some model definitions get remorselessly copy-pasted just because theforwardfunction does not return what the person expects. You provide module names
andtorchextractortakes care of the extraction for you.It's never been easier to
extract feature, add an extra loss or plug another head to a network.
Ler us know what amazing things you build withtorchextractor!Installationpipinstalltorchextractor# stablepipinstallgit+https://github.com/antoinebrl/torchextractor.git# latestRequirements:Python >= 3.6+torch >= 1.4.0Usageimporttorchimporttorchvisionimporttorchextractorastxmodel=torchvision.models.resnet18(pretrained=True)model=tx.Extractor(model,["layer1","layer2","layer3","layer4"])dummy_input=torch.rand(7,3,224,224)model_output,features=model(dummy_input)feature_shapes={name:f.shapeforname,finfeatures.items()}print(feature_shapes)# {#   'layer1': torch.Size([1, 64, 56, 56]),#   'layer2': torch.Size([1, 128, 28, 28]),#   'layer3': torch.Size([1, 256, 14, 14]),#   'layer4': torch.Size([1, 512, 7, 7]),# }See more examplesRead the documentationFAQ• How do I know the names of the modules?You can print all module names like this:tx.list_module_names(model)# ORforname,moduleinmodel.named_modules():print(name)• Why do some operations not get listed?It is not possible to add hooks if operations are not defined as modules.
Therefore,F.relucannot be captured butnn.Relu()can.• How can I avoid listing all relevant modules?You can specify a custom filtering function to hook the relevant modules:# Hook everything !module_filter_fn=lambdamodule,name:True# Capture of all modules inside first layermodule_filter_fn=lambdamodule,name:name.startswith("layer1")# Focus on all convolutionsmodule_filter_fn=lambdamodule,name:isinstance(module,torch.nn.Conv2d)model=tx.Extractor(model,module_filter_fn=module_filter_fn)• Is it compatible with ONNX?tx.Extractoris compatible with ONNX! This means you can also access intermediate features maps after the export.Pro-tip: name the output nodes by usingoutput_nameswhen callingtorch.onnx.export.• Is it compatible with TorchScript?Bad news, TorchScript cannot take variable number of arguments and keyword-only arguments.Good news, there is a workaround! The solution is to overwrite theforwardfunction
oftx.Extractorto replicate the interface of the model.importtorchimporttorchvisionimporttorchextractorastxclassMyExtractor(tx.Extractor):defforward(self,x1,x2,x3):# Assuming the model takes x1, x2 and x3 as inputoutput=self.model(x1,x2,x3)returnoutput,self.feature_mapsmodel=torchvision.models.resnet18(pretrained=True)model=MyExtractor(model,["layer1","layer2","layer3","layer4"])model_traced=torch.jit.script(model)• "One more thing!" :wink:By default we capture the latest output of the relevant modules,
but you can specify your own custom operations.For example, to accumulate features over 10 forward passes you
can do the following:importtorchimporttorchvisionimporttorchextractorastxmodel=torchvision.models.resnet18(pretrained=True)defcapture_fn(module,input,output,module_name,feature_maps):ifmodule_namenotinfeature_maps:feature_maps[module_name]=[]feature_maps[module_name].append(output)extractor=tx.Extractor(model,["layer3","layer4"],capture_fn=capture_fn)foriinrange(20):foriinrange(10):x=torch.rand(7,3,224,224)model(x)feature_maps=extractor.collect()# Do your stuffs here# Discard collected elementsextractor.clear_placeholder()ContributingAll feedbacks and contributions are welcomed. Feel free to report an issue or to create a pull request!If you want to get hands-on:(Fork and) clone the repo.Create a virtual environment:virtualenv -p python3 .venv && source .venv/bin/activateInstall dependencies:pip install -r requirements.txt && pip install -r requirements-dev.txtHook auto-formatting tools:pre-commit installHack as much as you want!Run tests:python -m unittest discover -vs ./tests/Share your work and create a pull request.To Build documentation:cddocs
pipinstallrequirements.txt
makehtml