Torch-Pruning 项目使用指南

Torch-Pruning 项目使用指南

Torch-Pruning[CVPR 2023] Towards Any Structural Pruning; LLMs / Diffusion / Transformers / YOLOv8 / CNNs项目地址:https://gitcode.com/gh_mirrors/to/Torch-Pruning

1. 项目的目录结构及介绍

Torch-Pruning 项目的目录结构如下：

Torch-Pruning/
├── AUTHORS
├── LICENSE
├── README.md
├── README_CN.md
├── practical_structural_pruning.md
├── requirements.txt
├── setup.py
├── github/
│   └── workflows/
├── assets/
├── examples/
├── reproduce/
│   └── main.py
├── tests/
└── torch_pruning/

目录介绍

• AUTHORS: 项目作者列表。
• LICENSE: 项目许可证文件。
• README.md: 项目英文介绍文档。
• README_CN.md: 项目中文介绍文档。
• practical_structural_pruning.md: 结构化剪枝实践文档。
• requirements.txt: 项目依赖文件。
• setup.py: 项目安装脚本。
• github/workflows/: GitHub Actions 工作流配置文件。
• assets/: 项目资源文件。
• examples/: 示例代码目录。
• reproduce/: 复现论文结果的代码目录。
• tests/: 测试代码目录。
• torch_pruning/: 项目核心代码目录。

2. 项目的启动文件介绍

项目的启动文件位于 reproduce/main.py。这个文件包含了结构化剪枝的示例代码，展示了如何使用 Torch-Pruning 库进行模型剪枝。

启动文件内容概述

import torch
from torchvision.models import resnet18
import torch_pruning as tp

# 加载预训练模型
model = resnet18(pretrained=True)

# 示例输入
example_inputs = torch.randn(1, 3, 224, 224)

# 重要性准则
imp = tp.importance.GroupNormImportance(p=2)

# 初始化剪枝器
ignored_layers = []
for m in model.modules():
    if isinstance(m, torch.nn.Linear) and m.out_features == 1000:
        ignored_layers.append(m)

pruner = tp.pruner.MetaPruner(
    model,
    example_inputs,
    importance=imp,
    pruning_ratio=0.5,  # 移除50%的通道
    ignored_layers=ignored_layers
)

# 执行剪枝
pruner.step()

3. 项目的配置文件介绍

项目中没有显式的配置文件，但可以通过 setup.py 和 requirements.txt 文件来管理项目的依赖和安装。

setup.py

setup.py 文件用于项目的安装和打包。它包含了项目的元数据和依赖信息。

from setuptools import setup, find_packages

setup(
    name='torch-pruning',
    version='1.0',
    packages=find_packages(),
    install_requires=[
        'torch',
        'torchvision'
    ],
    author='VainF',
    description='Structural pruning for PyTorch models',
    license='MIT',
    keywords='pytorch pruning',
    url='https://github.com/VainF/Torch-Pruning'
)

requirements.txt

requirements.txt 文件列出了项目运行所需的依赖包。

torch
torchvision

通过这两个文件，用户可以轻松地安装和管理项目的依赖。

Torch-Pruning[CVPR 2023] Towards Any Structural Pruning; LLMs / Diffusion / Transformers / YOLOv8 / CNNs项目地址:https://gitcode.com/gh_mirrors/to/Torch-Pruning