Pytorch:利用ResNet预训练模型对CIFAR数据集分类

电玩女神 2023-02-15 10:09 48阅读 0赞

> 最近看了常用的图像分类网络，分别使用迁移学习和直接构建模型自己训练。找一个模型把使用过程写下来吧。

### 文章目录 ###

*   *   *   *  1、加载数据集并预处理
             *  2、加载模型和预训练权重文件
             *  3. 修改网络模型
             *  4. 设置网络参数
             *  5. 训练
             *  6、测试

#### 1、加载数据集并预处理 ####

batch_size = 50
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    # data pre-treatment
    data_transform = { 
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
        "val": transforms.Compose([transforms.Resize(256),
                                   transforms.CenterCrop(224),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}
    
    # load train data
    trainset = torchvision.datasets.CIFAR10(root='data/cifar-10-python',
                                            train=True,
                                            download=False,
                                            transform=data_transform["train"])
    
    trainloader = Data.DataLoader(dataset=trainset,
                                  batch_size=batch_size,
                                  shuffle=True,
                                  num_workers=0)
    
    # load test data
    testset = torchvision.datasets.CIFAR10(root='data/cifar-10-python',
                                           train=False,
                                           download=False,
                                           transform=data_transform["val"])
    testloader = Data.DataLoader(dataset=testset,
                                 batch_size=batch_size,
                                 shuffle=False,
                                 num_workers=0)

#### 2、加载模型和预训练权重文件 ####

这里加载模型，可以有两种方式。  
**1.** 直接利用torchvision自带的模型进行下载并载入权重

# load model and load parameter into the model
    net = torchvision.models.resnet34(pretrained= True)
    net.to(device) # net into cuda

*缺点*：就是下载很缓慢，一般预训练权重文件都是几百兆，所以会等很久  
**2.** 自己提前下载好权重文件

net = torchvision.models.resnet34(pretrained= False)
    # load pretrain weights
    model_weight_path = "./weights/resnet34-pre.pth"
    net.load_state_dict(torch.load(model_weight_path), strict=False)
    net.to(device)  # net into cuda

只要把pretrained参数设置为False即可。再讲自己提前下载好的文件载入模型即可。（**tips**：迅雷下载很快）

如果想自己训练模型，不需要用预训练数据，则不用加载预训练文件即可，pytorch会自动进行参数初始化。

#### 3. 修改网络模型 ####

这里是resnet是使用ImageNet数据集进行训练，所以它的输出的\*\*num\_class = 1000\*，如下图：  
![在这里插入图片描述][20200604233209940.png]  
而对于自己的数据集，没有1000个类别，所以设置一下模型最后一层的输出。

# change fc layer structure
    inchannel = net.fc.in_features
    net.fc = nn.Linear(inchannel, 10)  # change the number of class

因为CIFAR数据集是十个类别，所以改为自己数据集的类别数即可。  
![在这里插入图片描述][20200604233424968.png]  
完成修改

#### 4. 设置网络参数 ####

这一层是考虑需不需要冻结模型参数，一般对于参数很庞大的模型，最好还是需要的，比如VGG网络，否则有可能会出现[CUDA out of memory][]的错误。

#对于模型的每个权重，使其不进行反向传播，即固定参数
    for param in net.parameters():
        param.requires_grad = False
    #但是参数全部固定了，也没法进行学习，所以我们不固定最后一层，即全连接层
    for param in net.fc.parameters():
        param.requires_grad = True

我建议在硬件条件允许的情况下，可以不使用这个固定参数的模型，对所有层的参数进行学习的精确性肯定更好，resnet本身这个模型的参数不是很庞大。

#### 5. 训练 ####

# define ooptimizer and loss function
    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(), lr=0.0001)
    # training
    best_acc = 0.0
    save_path = './weights/resNet34.pth'
    for epoch in range(1):
        # training
        net.train()
        running_loss = 0.0
        for step, data in enumerate(trainloader, start=0):
            images, labels = data
            optimizer.zero_grad()
            logits = net(images.to(device))
            loss = loss_function(logits, labels.to(device))
            loss.backward()
            optimizer.step()
    
            # print statistics
            running_loss += loss.item()
            # print train process
            rate = (step + 1) / len(trainloader)
            a = "*" * int(rate * 50)
            b = "." * int((1 - rate) * 50)
            print("\rtrain loss: {:^3.0f}%[{}->{}]{:.4f}".format(int(rate * 100), a, b, loss), end="")

#### 6、测试 ####

net.eval()  # change into test model
    acc = 0.0  # accumulate accurate number / epoch
    with torch.no_grad():
        for val_data in testloader:
            val_images, val_labels = val_data
            outputs = net(val_images.to(device))  # eval model only have last output layer
            # loss = loss_function(outputs, test_labels)
            predict_y = torch.max(outputs, dim=1)[1]
            acc += (predict_y == val_labels.to(device)).sum().item()
        val_accurate = acc / test_num
        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(net.state_dict(), save_path)
        print('[epoch %d] train_loss: %.3f test_accuracy: %.3f' %
              (epoch + 1, running_loss / step, val_accurate))

测试不要忘了将模型转为测试模式**net.eval()** 并且使用 **with torch.no\_grad()** 能够减少算力，并且不会保存梯度减少内存。

[20200604233209940.png]: https://img-blog.csdnimg.cn/20200604233209940.png
[20200604233424968.png]: https://img-blog.csdnimg.cn/20200604233424968.png
[CUDA out of memory]: https://blog.csdn.net/qq_35896136/article/details/106529986