{"id":13577814,"url":"https://github.com/antoinebrl/torchextractor","last_synced_at":"2025-08-20T04:32:28.871Z","repository":{"id":65463307,"uuid":"344262572","full_name":"antoinebrl/torchextractor","owner":"antoinebrl","description":"Feature extraction made simple with torchextractor","archived":false,"fork":false,"pushed_at":"2021-03-10T09:51:56.000Z","size":32,"stargazers_count":99,"open_issues_count":1,"forks_count":4,"subscribers_count":2,"default_branch":"main","last_synced_at":"2024-04-27T06:53:58.224Z","etag":null,"topics":["feature-extraction","machine-learning","python3","pytorch"],"latest_commit_sha":null,"homepage":"https://torchextractor.readthedocs.io","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/antoinebrl.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2021-03-03T21:01:16.000Z","updated_at":"2024-02-23T18:18:14.000Z","dependencies_parsed_at":null,"dependency_job_id":"c7303479-3d01-4134-a71e-45b0e539bf6b","html_url":"https://github.com/antoinebrl/torchextractor","commit_stats":null,"previous_names":[],"tags_count":2,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/antoinebrl%2Ftorchextractor","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/antoinebrl%2Ftorchextractor/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/antoinebrl%2Ftorchextractor/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/antoinebrl%2Ftorchextractor/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/antoinebrl","download_url":"https://codeload.github.com/antoinebrl/torchextractor/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":230394228,"owners_count":18218707,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["feature-extraction","machine-learning","python3","pytorch"],"created_at":"2024-08-01T15:01:24.562Z","updated_at":"2024-12-19T07:06:41.283Z","avatar_url":"https://github.com/antoinebrl.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"# `torchextractor`: PyTorch Intermediate Feature Extraction\n\n[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/torchextractor)](https://pypi.org/project/torchextractor/)\n[![PyPI](https://img.shields.io/pypi/v/torchextractor)](https://pypi.org/project/torchextractor/)\n[![Read the Docs](https://img.shields.io/readthedocs/torchextractor)](https://torchextractor.readthedocs.io/en/latest/)\n[![Upload Python Package](https://github.com/antoinebrl/torchextractor/actions/workflows/publish.yml/badge.svg)](https://github.com/antoinebrl/torchextractor/actions/workflows/publish.yml)\n[![GitHub](https://img.shields.io/github/license/antoinebrl/torchextractor)](https://github.com/antoinebrl/torchextractor/blob/main/LICENSE)\n    \n\n## Introduction\n\nToo many times some model definitions get remorselessly copy-pasted just because the\n`forward` function does not return what the person expects. You provide module names\nand `torchextractor` takes care of the extraction for you.It's never been easier to\nextract feature, add an extra loss or plug another head to a network.\nLer us know what amazing things you build with `torchextractor`!\n\n## Installation\n\n```shell\npip install torchextractor  # stable\npip install git+https://github.com/antoinebrl/torchextractor.git  # latest\n```\n\nRequirements:\n- Python \u003e= 3.6+\n- torch \u003e= 1.4.0\n\n## Usage\n\n```python\nimport torch\nimport torchvision\nimport torchextractor as tx\n\nmodel = torchvision.models.resnet18(pretrained=True)\nmodel = tx.Extractor(model, [\"layer1\", \"layer2\", \"layer3\", \"layer4\"])\ndummy_input = torch.rand(7, 3, 224, 224)\nmodel_output, features = model(dummy_input)\nfeature_shapes = {name: f.shape for name, f in features.items()}\nprint(feature_shapes)\n\n# {\n#   'layer1': torch.Size([1, 64, 56, 56]),\n#   'layer2': torch.Size([1, 128, 28, 28]),\n#   'layer3': torch.Size([1, 256, 14, 14]),\n#   'layer4': torch.Size([1, 512, 7, 7]),\n# }\n```\n\n[See more examples](docs/source/examples.ipynb)\n[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/antoinebrl/torchextractor/HEAD?filepath=docs/source/examples.ipynb)\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/antoinebrl/torchextractor/blob/master/docs/source/examples.ipynb)\n\n[Read the documentation](https://torchextractor.readthedocs.io/en/latest/)\n\n## FAQ\n\n**• How do I know the names of the modules?**\n\nYou can print all module names like this:\n```python\ntx.list_module_names(model)\n\n# OR\n\nfor name, module in model.named_modules():\n    print(name)\n```\n\n**• Why do some operations not get listed?**\n\nIt is not possible to add hooks if operations are not defined as modules.\nTherefore, `F.relu` cannot be captured but `nn.Relu()` can.\n\n**• How can I avoid listing all relevant modules?**\n\nYou can specify a custom filtering function to hook the relevant modules:\n```python\n# Hook everything !\nmodule_filter_fn = lambda module, name: True\n\n# Capture of all modules inside first layer\nmodule_filter_fn = lambda module, name: name.startswith(\"layer1\")\n\n# Focus on all convolutions\nmodule_filter_fn = lambda module, name: isinstance(module, torch.nn.Conv2d)\n\nmodel = tx.Extractor(model, module_filter_fn=module_filter_fn)\n```\n\n**• Is it compatible with ONNX?**\n\n`tx.Extractor` is compatible with ONNX! This means you can also access intermediate features maps after the export.\n\nPro-tip: name the output nodes by using `output_names` when calling `torch.onnx.export`.\n\n**• Is it compatible with TorchScript?**\n\nNot yet, but we are working on it. Compiling registered hook of a module\n[was just recently added in PyTorch v1.8.0](https://github.com/pytorch/pytorch/pull/49544).\n\n**• \"One more thing!\" :wink:**\n\nBy default we capture the latest output of the relevant modules,\nbut you can specify your own custom operations.\n\nFor example, to accumulate features over 10 forward passes you\ncan do the following:\n```python\nimport torch\nimport torchvision\nimport torchextractor as tx\n\nmodel = torchvision.models.resnet18(pretrained=True)\n\ndef capture_fn(module, input, output, module_name, feature_maps):\n    if module_name not in feature_maps:\n        feature_maps[module_name] = []\n    feature_maps[module_name].append(output)\n\nextractor = tx.Extractor(model, [\"layer3\", \"layer4\"], capture_fn=capture_fn)\n\nfor i in range(20):\n    for i in range(10):\n        x = torch.rand(7, 3, 224, 224)\n        model(x)\n    feature_maps = extractor.collect()\n\n    # Do your stuffs here\n\n    # Discard collected elements\n    extractor.clear_placeholder()\n```\n\n## Contributing\n\nAll feedbacks and contributions are welcomed. Feel free to report an issue or to create a pull request!\n\nIf you want to get hands-on:\n1. (Fork and) clone the repo.\n2. Create a virtual environment: `virtualenv -p python3 .venv \u0026\u0026 source .venv/bin/activate`\n2. Install dependencies: `pip install -r requirements.txt \u0026\u0026 pip install -r requirements-dev.txt`\n4. Hook auto-formatting tools: `pre-commit install`\n5. Hack as much as you want!\n6. Run tests: `python -m unittest discover -vs ./tests/`\n7. Share your work and create a pull request.\n\nTo Build documentation:\n```shell\ncd docs\npip install requirements.txt\nmake html\n```","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fantoinebrl%2Ftorchextractor","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fantoinebrl%2Ftorchextractor","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fantoinebrl%2Ftorchextractor/lists"}