pytorch超分辨率剪枝代码

以下是使用PyTorch实现的超分辨率剪枝代码的示例：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
from torchvision import transforms
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import numpy as np
import os
from math import log10

# 定义超分辨率网络
class SuperResolutionNet(nn.Module):
    def __init__(self):
        super(SuperResolutionNet, self).__init__()
        # 定义网络结构
        self.layer1 = nn.Sequential(nn.Conv2d(3, 64, (5, 5), (1, 1), (2, 2)),
                                     nn.ReLU())
        self.layer2 = nn.Sequential(nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
                                     nn.ReLU())
        self.layer3 = nn.Sequential(nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1)),
                                     nn.ReLU())
        self.layer4 = nn.Sequential(nn.Conv2d(32, 3, (3, 3), (1, 1), (1, 1)))

    def forward(self, x):
        # 前向传播
        out1 = self.layer1(x)
        out2 = self.layer2(out1)
        out3 = self.layer3(out2)
        out4 = self.layer4(out3)
        return out4

# 定义超分辨率数据集
class SuperResolutionDataset(data.Dataset):
    def __init__(self, image_folder, transform=None):
        super(SuperResolutionDataset, self).__init__()
        # 加载图像文件
        self.image_folder = image_folder
        self.image_filenames = [os.path.join(self.image_folder, x)
                                for x in os.listdir(self.image_folder)
                                if is_image_file(x)]
        self.transform = transform

    def __getitem__(self, index):
        # 获取图像和目标
        input = load_img(self.image_filenames[index])
        target = input.copy()
        # 转换图像
        if self.transform:
            input = self.transform(input)
            target = self.transform(target)
        # 返回输入和目标
        return input, target

    def __len__(self):
        # 获取数据集大小
        return len(self.image_filenames)

# 定义图片载入函数
def load_img(filepath):
    img = Image.open(filepath).convert('RGB')
    img = np.array(img).astype(np.float32)
    img = img / 255.0
    return img

# 定义图片类型判断函数
def is_image_file(filename):
    return any(filename.endswith(extension)
               for extension in ['.png', '.jpg', '.jpeg'])

# 定义超分辨率训练函数
def train(epoch):
    model.train()
    train_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        if cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        train_loss += loss.data[0]
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader),
                       loss.data[0] / len(data)))
    print('====> Epoch: {} Average loss: {:.4f}'.format(
        epoch, train_loss / len(train_loader.dataset)))

# 定义超分辨率测试函数
def test(epoch):
    model.eval()
    test_loss = 0
    for batch_idx, (data, target) in enumerate(test_loader):
        if cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data, volatile=True), Variable(target)
        output = model(data)
        test_loss += criterion(output, target).data[0]
        psnr = 10 * log10(1 / test_loss)
        if batch_idx % log_interval == 0:
            print('Test Epoch: {} [{}/{} ({:.0f}%)]\tPSNR: {:.6f}'.format(
                epoch, batch_idx * len(data), len(test_loader.dataset),
                       100. * batch_idx / len(test_loader),
                       psnr))
    print('====> Epoch: {} Average PSNR: {:.4f}'.format(
        epoch, psnr))

# 定义超分辨率剪枝函数
def prune(model, pruning_perc):
    # 获取模型权重
    weights = []
    for name, param in model.named_parameters():
        if 'weight' in name:
            weights.append(param.data.cpu().numpy().flatten())
    weights = np.concatenate(weights)
    # 计算权重阈值
    threshold = np.percentile(abs(weights), pruning_perc)
    # 定义剪枝函数
    def prune_weights(weights, threshold):
        return np.where(abs(weights) > threshold, weights, 0)
    # 剪枝模型权重
    for name, param in model.named_parameters():
        if 'weight' in name:
            w = param.data.cpu().numpy()
            param.data = torch.from_numpy(prune_weights(w, threshold)).cuda()

# 加载超分辨率数据集
train_dataset = SuperResolutionDataset(image_folder='train', transform=transforms.ToTensor())
test_dataset = SuperResolutionDataset(image_folder='test', transform=transforms.ToTensor())
# 定义超分辨率数据加载器
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=True)
# 定义超分辨率网络
model = SuperResolutionNet()
# 定义超分辨率损失函数
criterion = nn.MSELoss()
# 定义超分辨率优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 定义超分辨率训练参数
epochs = 10
log_interval = 10
cuda = True
# 训练超分辨率网络
for epoch in range(1, epochs + 1):
    train(epoch)
    test(epoch)
    # 每个 epoch 结束后对模型进行剪枝
    prune(model, 50)