RepLKNet-pytorch (CVPR 2022)

这是 RepLKNet 的官方 PyTorch 实现，源自以下 CVPR-2022 论文：

将卷积核扩大到 31x31：重新审视 CNN 中的大卷积核设计。

该论文现已发布在 arXiv 上：https://arxiv.org/abs/2203.06717。

更新：所有预训练模型、ImageNet-1K 模型以及 Cityscapes/ADE20K/COCO 模型均已发布。

更新：发布了用于可视化有效感受野（ERF）的脚本。只需添加几行代码，即可获取您自己模型的有效感受野！

更新：发布了训练命令及更多示例。

如果您觉得这篇论文或本仓库对您有所帮助，请考虑引用：

    @article{replknet,
    title={Scaling Up Your Kernels to 31x31: Revisiting Large Kernel Design in CNNs},
    author={Ding, Xiaohan and Zhang, Xiangyu and Zhou, Yizhuang and Han, Jungong and Ding, Guiguang and Sun, Jian},
    journal={arXiv preprint arXiv:2203.06717},
    year={2022}
    }

其他实现

框架	链接
MegEngine（官方）	https://github.com/megvii-research/RepLKNet
PyTorch（官方）	https://github.com/DingXiaoH/RepLKNet-pytorch
TensorFlow	https://github.com/shkarupa-alex/tfreplknet
PaddlePaddle	https://github.com/BR-IDL/PaddleViT/tree/develop/image_classification/RepLKNet
...

欢迎更多复现版本。

使用我们的高效大卷积核卷积（PyTorch）

我们已为 PyTorch 发布了一个示例。请查看 https://github.com/MegEngine/cutlass/tree/master/examples/19_large_depthwise_conv2d_torch_extension 中的 setup.py 和 depthwise_conv2d_implicit_gemm.py（替代 torch.nn.Conv2d）。

1. 解压 cutlass.zip 并进入目录。此版本的 cutlass 与我们的大卷积核实现以及多个 Python 版本兼容。您也可以使用 MegEngine 团队维护的 cutlass 分支（克隆 https://github.com/MegEngine/cutlass），但可能需要更注意您的 Python 版本（参见 此问题）。
2. 进入 cd examples/19_large_depthwise_conv2d_torch_extension
3. 运行 ./setup.py install --user。如果出现错误，请检查您的 CUDA_HOME。
4. 快速测试：python depthwise_conv2d_implicit_gemm.py
5. 将 WHERE_YOU_CLONED_CUTLASS/examples/19_large_depthwise_conv2d_torch_extension 添加到您的 PYTHONPATH 中，以便您可以在任何地方 from depthwise_conv2d_implicit_gemm import DepthWiseConv2dImplicitGEMM。然后您可以使用 DepthWiseConv2dImplicitGEMM 替代 nn.Conv2d。
6. 设置环境变量 export LARGE_KERNEL_CONV_IMPL=WHERE_YOU_CLONED_CUTLASS/examples/19_large_depthwise_conv2d_torch_extension，这样 RepLKNet 就会使用高效的实现。或者您也可以直接修改 replknet.py 中的相关代码（get_conv2d）。

它应该能在多种 GPU 和 PyTorch/CUDA 版本上正常工作。我们建议您先尝试，只有在遇到错误时再检查环境。我们最新的测试使用了以下配置：

1. Ubuntu 18.04 + CUDA 11.3 + nvcc 11.3 + cudnn 8.2.0 + python 3.8.12 + pytorch 1.10 + gcc 7.3.0 + nccl 2.10.3 + NVIDIA 驱动程序 450.102.04 + V100 和 A100 GPU
2. Ubuntu 18.04 + CUDA 10.2 + nvcc 10.0 + cudnn 7.6.5 + python 3.6.9 + pytorch 1.9 + gcc 7.5.0 + nccl 2.7.8 + NVIDIA 驱动程序 460.32.03 + 2080Ti 和 V100 GPU

据报道（参见 这里），Python 版本不匹配可能会导致错误（forward_fp32.cu(212): error: more than one instance of constructor "cutlass::Tensor4DCoord::Tensor4DCoord" ... 或 cutlass/include/cutlass/fast_math.h(741): error: no suitable conversion function from "__half" to "float" exists)。请升级或降级您的 Python 版本。我们衷心感谢 @sleeplessai 和 @ewrfcas 分享他们的经验。

我们在论文中提到的实现已被集成到 MegEngine 中，引擎会自动使用它。如果您想在其他框架如 TensorFlow 中使用，可能需要编译我们发布的 CUDA 源代码（上述示例中的 *.cu 文件应适用于其他框架），并使用一些工具来加载它们，就像 PyTorch 示例中的 cutlass 和 torch.utils.cpp_extension 一样。如果您能与我们分享您的经验，我们将不胜感激。

您可以参考 MegEngine 的源代码：https://github.com/MegEngine/MegEngine/tree/8a2e92bd6c5ac02807b27d174dce090ee391000b/dnn/src/cuda/conv_bias/chanwise。

欢迎提交 Pull 请求（例如更好的实现、其他实现或针对其他框架的实现）。

目录

• 模型代码
• PyTorch 预训练模型
• PyTorch 大卷积核卷积实现
• PyTorch 训练代码
• PyTorch 下游模型
• PyTorch 下游代码
• 用于可视化 ERF 的脚本
• 如何获取形状偏差

结果与预训练模型

ImageNet-1K 模型

名称	分辨率	ImageNet-1K 准确率	参数量	FLOPs	ImageNet-1K 预训练模型
RepLKNet-31B	224x224	83.5	79M	15.3G	Google Drive, 百度网盘
RepLKNet-31B	384x384	84.8	79M	45.1G	Google Drive, 百度网盘

ImageNet-22K 模型

名称	分辨率	ImageNet-1K 准确率	参数量	FLOPs	22K 预训练模型	1K 微调后模型
RepLKNet-31B	224x224	85.2	79M	15.3G	Google Drive, 百度网盘	Google Drive, 百度网盘
RepLKNet-31B	384x384	86.0	79M	45.1G	-	Google Drive, 百度网盘
RepLKNet-31L	384x384	86.6	172M	96.0G	Google Drive, 百度网盘	Google Drive, 百度网盘

MegData-73M 模型

名称	分辨率	ImageNet-1K 准确率	参数量	FLOPs	MegData-73M 预训练模型	1K 微调模型
RepLKNet-XL	320x320	87.8	335M	128.7G	Google Drive, Baidu	Google Drive, Baidu

评估

对于分辨率为 224x224 或 384x384 的 RepLKNet-31B/L，我们使用 "IMAGENET_DEFAULT_MEAN/STD" 进行预处理（参见 此处）。例如：

python -m torch.distributed.launch --nproc_per_node=8 main.py --model RepLKNet-31B --batch_size 32 --eval True --resume RepLKNet-31B_ImageNet-1K_224.pth --input_size 224

或者

python -m torch.distributed.launch --nproc_per_node=8 main.py --model RepLKNet-31L --batch_size 32 --eval True --resume RepLKNet-31L_ImageNet-22K-to-1K_384.pth --input_size 384

对于 RepLKNet-XL，请注意，我们在 MegData73M 数据集以及 ImageNet-1K 上的微调中，均使用了 mean=[0.5,0.5,0.5] 和 std=[0.5,0.5,0.5] 进行预处理。这种均值和标准差设置在 timm 中也被称为 "IMAGENET_INCEPTION_MEAN/STD"，详见 此处。要使用此均值和标准差设置，请添加 --imagenet_default_mean_and_std false（参见 此处）。正如论文中所提到的，我们并未使用小卷积核进行重参数化。

python -m torch.distributed.launch --nproc_per_node=8 main.py --model RepLKNet-XL --batch_size 32 --eval true --resume RepLKNet-XL_MegData73M_ImageNet1K.pth --imagenet_default_mean_and_std false --input_size 320

为了验证结构重参数化的等价性（即在执行 structural_reparam 之前和之后的输出是否相同），请添加 --with_small_kernel_merged true。

训练

您可以在 SLURM 集群上使用 submitit 进行多节点训练。请先安装：

pip install submitit

如果您 GPU 显存有限（例如 2080Ti），可以使用 --use_checkpoint true 来节省显存。

在 ImageNet-1K 上预训练 RepLKNet-31B

单机（注意 --update_freq 4）：

python -m torch.distributed.launch --nproc_per_node=8 main.py --model RepLKNet-31B --drop_path 0.5 --batch_size 64 --lr 4e-3 --update_freq 4 --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 10 --epochs 300 --output_dir your_training_dir

四台机器（注意 --update_freq 1）：

python run_with_submitit.py --nodes 4 --ngpus 8 --model RepLKNet-31B --drop_path 0.5 --batch_size 64 --lr 4e-3 --update_freq 1 --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 10 --epochs 300 --job_dir your_training_dir

以下仅列出多机命令。您也可以用类似的方式在单机上进行训练。

将 ImageNet-1K 预训练的（224x224）RepLKNet-31B 微调至 384x384

python run_with_submitit.py --nodes 4 --ngpus 8 --model RepLKNet-31B --drop_path 0.8 --input_size 384 --batch_size 32 --lr 4e-4 --epochs 30 --weight_decay 1e-8 --update_freq 1 --cutmix 0 --mixup 0 --finetune RepLKNet-31B_ImageNet-1K_224.pth --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 1 --job_dir your_training_dir --layer_decay 0.7

在 ImageNet-22K 上预训练 RepLKNet-31B

python run_with_submitit.py --nodes 16 --ngpus 8 --model RepLKNet-31B --drop_path 0.1 --batch_size 32 --lr 4e-3 --update_freq 1 --warmup_epochs 5 --epochs 90 --data_set image_folder --nb_classes 21841 --disable_eval true --data_path /path/to/imagenet-22k --job_dir /path/to/save_results

将 22K 预训练的 RepLKNet-31B 微调至 ImageNet-1K（224x224）

python run_with_submitit.py --nodes 2 --ngpus 8 --model RepLKNet-31B --drop_path 0.2 --input_size 224 --batch_size 32 --lr 4e-4 --epochs 30 --weight_decay 1e-8 --update_freq 1 --cutmix 0 --mixup 0 --finetune RepLKNet-31B_ImageNet-22K.pth --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 1 --job_dir your_training_dir --layer_decay 0.7

将 22K 预训练的 RepLKNet-31B 微调至 ImageNet-1K（384x384）

python run_with_submitit.py --nodes 4 --ngpus 8 --model RepLKNet-31B --drop_path 0.3 --input_size 384 --batch_size 16 --lr 4e-4 --epochs 30 --weight_decay 1e-8 --update_freq 1 --cutmix 0 --mixup 0 --finetune RepLKNet-31B_ImageNet-22K.pth --model_ema true --model_ema_eval true --data_path /path/to/imagenet-1k --warmup_epochs 1 --job dir your training dir --layer decay 0.7 --min_lr 3e-4

在 ImageNet-22K 上预训练 RepLKNet-31L

python run_with_submitit.py --nodes 16 --ngpus 8 --model RepLKNet-31L --drop_path 0.1 --batch_size 32 --lr 4e-3 --update_freq 1 --warmup epochs 5 --epochs 90 --data_set image_folder --nb_classes 21841 --disable eval true --data path /path/to/imagenet-22k --job dir /path/to/save results

将 22K 预训练的 RepLKNet-31L 微调至 ImageNet-1K（384x384）

python run_with_submitit.py --nodes 4 --ngpus 8 --model RepLKNet-31L --drop_path 0.3 --input size 384 --batch size 16 --lr 4e-4 --epochs 30 --weight decay 1e-8 --update freq 1 --cutmix 0 --mixup 0 --finetune RepLKNet-31L_ImageNet-22K.pth --model ema true --model ema eval true --data path /path/to/imagenet-1k --warmup epochs 1 --job dir your training dir --layer decay 0.7 --min lr 3e-4

语义分割与目标检测

我们使用 MMSegmentation 和 MMDetection 框架。只需克隆 MMSegmentation 或 MMDetection，然后：

1. 将 segmentation/replknet.py 放入 mmsegmentation/mmseg/models/backbones/ 或 mmdetection/mmdet/models/backbones/ 目录中。segmentation/replknet.py 和 replknet.py 的唯一区别在于装饰器 @BACKBONES.register_module。
2. 在 mmsegmentation/mmseg/models/backbones/__init__.py 或 mmdetection/mmdet/models/backbones/__init__.py 中添加 RepLKNet。具体如下：

   ...
   from .replknet import RepLKNet
   __all__ = ['ResNet', ..., 'RepLKNet']
   

3. 将 segmentation/configs/*.py 文件放入 mmsegmentation/configs/replknet/ 目录，或将 detection/configs/*.py 文件放入 mmdetection/configs/replknet/ 目录。
4. 下载并使用我们的预训练权重。例如，在 Cityscapes 数据集上评估 RepLKNet-31B + UperNet：

   python -m torch.distributed.launch --nproc_per_node=8 tools/test.py configs/replknet/RepLKNet-31B_1Kpretrain_upernet_80k_cityscapes_769.py RepLKNet-31B_ImageNet-1K_UperNet_Cityscapes.pth --launcher pytorch --eval mIoU
   

   或者在 COCO 数据集上评估 RepLKNet-31B + Cascade Mask R-CNN：

   python -m torch.distributed.launch --nproc_per_node=8 tools/test.py configs/replknet/RepLKNet-31B_22Kpretrain_cascade_mask_rcnn_3x_coco.py RepLKNet-31B_ImageNet-22K_CascMaskRCNN_COCO.pth --eval bbox --launcher pytorch
   

5. 你也可以对已发布的预训练权重进行微调（有关批次大小和迭代次数的提示见下文）：

   python -m torch.distributed.launch --nproc_per_node=8 tools/train.py configs/replknet/some_config.py --launcher pytorch --options model.backbone.pretrained=some_pretrained_weights.pth
   

我们已发布了所有 Cityscapes/ADE20K/COCO 数据集的模型权重。

采用单尺度 (ss) 和多尺度 (ms) 测试的 mIoU 均基于 UperNet，FLOPs 计算时，ImageNet-1K 预训练模型使用 2048×512 分辨率，而 22K 和 MegData73M 预训练模型则使用 2560×640 分辨率，遵循 Swin 的做法：

主干网络	预训练数据集	数据集	训练计划	mIoU (ss)	mIoU (ms)	参数量	FLOPs	下载链接
RepLKNet-31B	ImageNet-1K	Cityscapes	8万次迭代	83.1	83.5	1.1亿	2315G	Google Drive, 百度网盘
RepLKNet-31B	ImageNet-1K	ADE20K	16万次迭代	49.9	50.6	1.12亿	1170G	Google Drive, 百度网盘
RepLKNet-31B	ImageNet-22K	ADE20K	16万次迭代	51.5	52.3	1.12亿	1829G	Google Drive, 百度网盘
RepLKNet-31L	ImageNet-22K	ADE20K	16万次迭代	52.4	52.7	2.07亿	2404G	Google Drive, 百度网盘
RepLKNet-XL	MegData73M	ADE20K	16万次迭代	55.2	56.0	3.74亿	3431G	Google Drive, 百度网盘

COCO 数据集上的 Cascade Mask R-CNN（FLOPs 计算时使用 1280×800 分辨率）：

主干网络	预训练数据集	方法	训练计划	AP_box	AP_mask	参数量	FLOPs	下载链接
RepLKNet-31B	ImageNet-1K	FCOS	2倍周期	47.0	-	8700万	437G	Google Drive, 百度网盘
RepLKNet-31B	ImageNet-1K	Cascade Mask R-CNN	3倍周期	52.2	45.2	1.37亿	965G	Google Drive, 百度网盘
RepLKNet-31B	ImageNet-22K	Cascade Mask R-CNN	3倍周期	53.0	46.0	1.37亿	965G	Google Drive, 百度网盘
RepLKNet-31L	ImageNet-22K	Cascade Mask R-CNN	3倍周期	53.9	46.5	2.29亿	1321G	Google Drive, 百度网盘
RepLKNet-XL	MegData73M	Cascade Mask R-CNN	3倍周期	55.5	48.0	3.92亿	1958G	Google Drive, 百度网盘

预训练或微调的几点建议

1. MegData73M 数据集上的均值/标准差与 ImageNet 不同。因此，我们在 MegData73M 上对 RepLKNet-XL 进行预训练，并在 ImageNet-1K 上进行微调时，使用了 mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]。相应地，在 MMSegmentation 和 MMDetection 中，应设置 img_norm_cfg = dict(mean=[127.5, 127.5, 127.5], std=[127.5, 127.5, 127.5], to_rgb=True)。请参考 此处 和 此处。对于其他模型，我们则沿用默认的 ImageNet 均值/标准差。

2. 对于在 ADE20K 和 COCO 数据集上使用的 RepLKNet-XL，我们在将中间特征图输入到头部之前进行了批量归一化处理。只需使用 RepLKNet(..., norm_intermediate_features=True) 即可。我们尚未在其他模型上尝试这种设计，因此不确定其效果是否显著。

3. 对于在 Cityscapes 和 ADE20K 数据集上使用的 RepLKNet-31B/L，我们分别使用了 4 或 8 个配备 8 张 2080Ti 显卡的节点，每张显卡的批大小低于默认值（默认为每张显卡 4 个样本，详见 此处），但全局批大小更大。为此，我们减少了迭代次数，以确保总的训练样本数不变。请参阅配置文件中的注释。如果您希望使用我们的配置文件进行训练，请根据自身情况调整批大小和迭代次数。

4. 在 ImageNet-1K 或下游任务上进行微调时，降低低层网络的学习率可能会提升性能，这一点与 ConvNeXt 和 BeiT 类似。不过，我们尚不确定这种改进是否显著。对于 ImageNet，我们的实现直接参考了 ConvNeXt 和 BeiT 的做法。若在 MMSegmentation 和 MMDetection 中需要相关示例，请提交问题。

5. 关于 drop_path_rate 的建议：模型越大，drop_path 应越高；预训练数据越多，drop_path 则越低。

可视化有效感受野

我们已发布用于可视化和分析有效感受野（ERF）的脚本。例如，要自动下载 torchvision 中的 ResNet-101 模型并获取聚合贡献度矩阵，可以运行以下命令：

python erf/visualize_erf.py --model resnet101 --data_path /path/to/imagenet-1k --save_path resnet101_erf_matrix.npy

随后，可通过以下命令计算高贡献区域比例并可视化 ERF：

python erf/analyze_erf.py --source resnet101_erf_matrix.npy --heatmap_save resnet101_heatmap.png

请注意，该绘图脚本适用于 matplotlib 3.3 版本。若您使用更高版本的 matplotlib，请参考 此处 的注释说明。

若要可视化您自己的模型，首先需定义一个输出最后一层特征图而非 logits 的模型（参考 此示例），并在 此处 添加构建模型和加载权重的相关代码，然后运行：

python erf/visualize_erf.py --model your_model --weights /path/to/your/weights --data_path /path/to/imagenet-1k --save_path your_model_erf_matrix.npy

若想复现论文中的结果，请下载经过 120 个 epoch 训练的 RepLKNet-13（Google Drive，百度网盘）和 RepLKNet-31（Google Drive，百度网盘）模型。

如何获取形状偏差

1. 安装 https://github.com/bethgelab/model-vs-human。
2. 将您构建模型和加载权重的代码添加到 此文件 中。例如：

@register_model("pytorch")
def replknet(model_name, *args):
    model = ...
    model.load_state_dict(...)
    return model

3. 根据其示例，修改 examples/evaluate.py（models = ['replknet']）和 examples/plotting_definition.py（decision_makers.append(DecisionMaker(name_pattern="replknet", ...))）。

致谢

本次发布的 PyTorch 训练脚本基于 ConvNeXt 的代码，该代码使用了 timm 库、DeiT 和 BEiT 等项目的技术成果。

许可证

本项目采用 MIT 许可证开源。更多信息请参阅 LICENSE 文件。

联系方式

xiaohding@gmail.com（原清华大学邮箱 dxh17@mails.tsinghua.edu.cn 将于数月后失效）

Google Scholar 个人主页：https://scholar.google.com/citations?user=CIjw0KoAAAAJ&hl=en

个人主页：https://dingxiaohan.xyz/

我的开源论文与代码库：

结构重参数化宇宙：

1. RepLKNet（CVPR 2022）具有超大卷积核（31×31）的强大高效架构，并提供了在 CNN 模型中使用大卷积核的指导原则
   将卷积核扩大到 31×31：重新审视 CNN 中的大卷积核设计
   代码。

2. RepOptimizer 使用梯度重参数化高效训练强大模型。其训练时的模型结构与推理时完全一致，同时解决了量化问题。
   与其重参数化架构，不如重参数化优化器
   代码。

3. RepVGG（CVPR 2021）一种超级简单且强大的 VGG 风格卷积神经网络架构。在 ImageNet 上最高可达 84.16% 的 top-1 准确率！
   RepVGG：让 VGG 风格的卷积神经网络再次伟大
   代码。

4. RepMLP（CVPR 2022）MLP 风格的构建模块及架构
   RepMLPNet：具有重参数化局部性的层次化视觉 MLP
   代码。

5. ResRep（ICCV 2021）最先进的通道剪枝技术（ResNet50，FLOPs 减少 55%，准确率 76.15%）
   ResRep：通过解耦记忆与遗忘实现无损 CNN 剪枝
   代码。

6. ACB（ICCV 2019）是一种在推理阶段无需额外开销的 CNN 组件，也是我们“结构重参数化宇宙”的开山之作。
   ACNet：通过非对称卷积块强化卷积核骨架，打造更强大的 CNN
   代码。

7. DBB（CVPR 2021）是一种性能优于 ACB 且同样无需推理开销的 CNN 组件。有时我称它为 ACNet v2，因为“DBB”在 ASCII 码中比“ACB”多两位（哈哈）。
   多样化分支模块：将卷积层构建为类似 Inception 的单元
   代码。

模型压缩与加速：

1. （CVPR 2019）通道剪枝：向心式 SGD 用于剪枝结构复杂的超深卷积神经网络
   代码。

2. （ICML 2019）通道剪枝：近似 Oracle 过滤器剪枝用于破坏性 CNN 宽度优化
   代码。

3. （NeurIPS 2019）非结构化剪枝：全局稀疏动量 SGD 用于剪枝超深神经网络
   代码

RepLKNet-pytorch 快速上手指南

RepLKNet 是 CVPR 2022 提出的基于大卷积核（最大 31x31）的 CNN 架构。本指南帮助开发者快速在 PyTorch 环境中部署和使用该模型。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 18.04+)
• GPU: NVIDIA GPU (支持 V100, A100, 2080Ti 等)
• CUDA: 建议 CUDA 10.2 或 11.3+
• Python: 建议 Python 3.6.9 - 3.8.12 (版本不匹配可能导致编译错误)
• PyTorch: 建议 1.9 或 1.10+

前置依赖

确保已安装以下基础库：

pip install torch torchvision timm submitit

注意：submitit 用于多节点训练，单机训练可不安装。

安装步骤

1. 克隆代码库

git clone https://github.com/DingXiaoH/RepLKNet-pytorch.git
cd RepLKNet-pytorch

2. 配置高效大核卷积算子 (可选但推荐)

为了获得论文中所述的最佳性能，需编译自定义 CUDA 算子 (DepthWiseConv2dImplicitGEMM)。若跳过此步，代码将自动回退到标准实现，但速度较慢。

步骤 A: 获取 CUTLASS 源码 下载并解压项目提供的 cutlass 包，或从 MegEngine 仓库克隆：

# 方式一：使用项目提供的 zip (如有)
unzip cutlass.zip
# 方式二：从官方源克隆 (需注意 Python 版本兼容性)
git clone https://github.com/MegEngine/cutlass.git

步骤 B: 编译扩展 进入指定目录并安装：

cd cutlass/examples/19_large_depthwise_conv2d_torch_extension
python setup.py install --user

若报错，请检查 CUDA_HOME 环境变量是否正确指向 CUDA 安装路径。

步骤 C: 配置环境变量 将编译好的路径加入 PYTHONPATH，并设置环境变量让 RepLKNet 调用该实现：

# 请将路径替换为你实际的克隆路径
export PYTHONPATH=$PYTHONPATH:/path/to/cutlass/examples/19_large_depthwise_conv2d_torch_extension
export LARGE_KERNEL_CONV_IMPL=/path/to/cutlass/examples/19_large_depthwise_conv2d_torch_extension

3. 验证安装

运行快速检查脚本：

python depthwise_conv2d_implicit_gemm.py

若无报错，则环境配置成功。

基本使用

加载预训练模型进行推理

以下示例展示如何加载 ImageNet-1K 预训练的 RepLKNet-31B 模型并进行评估。

单卡/单机评估命令：

python -m torch.distributed.launch --nproc_per_node=8 main.py \
--model RepLKNet-31B \
--batch_size 32 \
--eval True \
--resume RepLKNet-31B_ImageNet-1K_224.pth \
--input_size 224 \
--data_path /path/to/imagenet-1k

关键参数说明：

• --model: 模型名称 (如 RepLKNet-31B, RepLKNet-31L, RepLKNet-XL)。
• --resume: 预训练权重文件路径。
  • 国内下载加速: 百度网盘提取码均为 lknt (详见项目 README 中的表格)。
• --input_size: 输入分辨率 (224, 320, 或 384)。
• --imagenet_default_mean_and_std:
  • 对于 RepLKNet-XL (MegData-73M 预训练)，必须添加 --imagenet_default_mean_and_std false。
  • 其他模型默认使用 ImageNet 标准均值方差，无需额外参数。

代码调用示例 (Python)

在您的 Python 脚本中直接调用：

import torch
from replknet import create_model

# 创建模型
model = create_model(
    model_name='repLKNet31B', 
    pretrained=True, 
    num_classes=1000
)
model.eval()

# 准备输入 (Batch, Channels, Height, Width)
input_tensor = torch.randn(1, 3, 224, 224)

# 前向传播
with torch.no_grad():
    output = model(input_tensor)

print(output.shape)

显存优化

如果您的 GPU 显存较小（如 2080Ti），在训练或大分辨率推理时，可启用梯度检查点以节省显存：

# 在训练命令中添加
--use_checkpoint true