pytorch1.4.0实现对超分辨率模型的剪枝代码

这里给出一个基于PyTorch 1.4.0实现对超分辨率模型的剪枝的示例代码，以SRResNet为例：

import torch
import torch.nn.utils.prune as prune
from models import SRResNet

# 加载模型
model = SRResNet()

# 定义剪枝策略
pruning_plan = prune.L1Unstructured(
    amount=0.2,
    name='weight'
)

# 针对模型中的某一层进行剪枝
module = model.conv1
pruning_plan.apply(module, name='weight')

# 查看剪枝后模型的大小
print("Sparsity in conv1.weight: {:.2f}%".format(
    100. * float(torch.sum(module.weight == 0))
    / float(module.weight.nelement())
))
print("Size of conv1.weight: {:.2f} KB".format(
    float(module.weight.nelement() * 4) / 1024
))

# 对整个模型进行剪枝
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        pruning_plan.apply(module, name='weight')

# 查看整个模型的大小
print("Size of pruned model: {:.2f} KB".format(
    float(sum([param.nelement() for name, param in model.named_parameters()]))
    * 4 / 1024
))

# 保存剪枝后的模型
torch.save(model.state_dict(), 'pruned_srresnet.pth')

需要注意的是，这里的SRResNet是自定义的超分辨率模型，如果要使用其他的超分辨率模型，需要进行相应的修改。同时，需要在代码开头导入相应的模型。