CIFAR10_ResNet

import torch
import torchvision.datasets
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
import torch.nn as nn

import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

# 调用torchvision中的models中的resnet网络结构
import torchvision.models as models

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

cuda_is_available = torch.cuda.is_available()

transform_train = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transforms_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# 超参数
Download = True
EPOCHS = 10
LR = 0.0001
BATCHSIZE = 64

train_data = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=transform_train, download=Download)
test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=transforms_test, download=Download)

train_size = len(train_data)
test_size = len(test_data)

train_dataloader = DataLoader(dataset=train_data, batch_size=BATCHSIZE, shuffle=True)
test_dataloader = DataLoader(dataset=test_data, batch_size=BATCHSIZE, shuffle=True)


# 调整已经训练好的ResNet网络

class ResNet(nn.Module):
    def __init__(self, num_classes=10):  # num_classes，此处为 二分类值为2
        super(ResNet, self).__init__()
        net = models.resnet18(pretrained=True)  # 从预训练模型加载resnet18网络参数
        net.classifier = nn.Sequential()  # 将分类层置空，下面将改变我们的分类层
        # self.features = net  # 保留resnet的特征层
        self.layer0 = nn.Sequential(
            net.conv1,
            net.bn1,
            net.relu,
            net.maxpool
        )
        self.layer1 = net.layer1
        self.layer2 = net.layer2
        self.layer3 = net.layer3
        self.layer4 = net.layer4
        self.layer5 = nn.Sequential(
            nn.Conv2d(512, 512, (3, 3), (2, 2), 1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.Conv2d(512, 512, (3, 3), (1, 1), 1),
            nn.BatchNorm2d(512),

            nn.Conv2d(512, 512, (3, 3), (1, 1), 1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.Conv2d(512, 1024, (3, 3), (1, 1), 1),
            nn.BatchNorm2d(1024),

        )
        self.avgpool = net.avgpool
        self.fc = nn.Linear(1024, 1000)

        self.classifier = nn.Sequential(  # 定义自己的分类层
            nn.Linear(1000, 64),  # 1000不能改变 ，由resnet18网络决定的，第二个参数为神经元个数可以微调
            nn.ReLU(True),
            nn.Dropout(0.5),
            nn.Linear(64, num_classes),
        )

    def forward(self, x):
        x = self.layer0(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = self.avgpool(x)

        x = x.view(x.size(0), -1)
        x = self.fc(x)
        x = self.classifier(x)
        return x


import ResNet_bySelf

# ResNet18 = ResNet_bySelf.ResNet18(3,10)

ResNet = ResNet()
print(ResNet)
from collections.abc import Iterable


def set_freeze_by_names(model, layer_names, freeze=True):
    if not isinstance(layer_names, Iterable):
        layer_names = [layer_names]
    for name, child in model.named_children():
        if name not in layer_names:
            continue
        for param in child.parameters():
            param.requires_grad = not freeze


def freeze_by_names(model, layer_names):
    set_freeze_by_names(model, layer_names, True)


def unfreeze_by_names(model, layer_names):
    set_freeze_by_names(model, layer_names, False)


# 先把features层冻住，只训练FC层
unfreeze_by_names(ResNet, ("features"))

# 优化器选择Adam
optimizer = torch.optim.Adam(ResNet.parameters(), lr=LR)
# 损失函数
loss_func = nn.CrossEntropyLoss()  # 目标标签是one-hotted

if cuda_is_available:
    loss_func.cuda()
    ResNet.cuda()
# 开始训练


train_loss = []
train_acc = []
test_loss = []
test_acc = []
for epoch in range(EPOCHS):
    print("第{}轮训练开始:".format(epoch + 1))
    total_train_loss = 0
    total_test_loss = 0
    total_train_acc = 0
    total_test_acc = 0
    for step, (b_x, b_y) in enumerate(train_dataloader):  # 分配batch data
        if cuda_is_available:
            b_x = b_x.cuda()
            b_y = b_y.cuda()
        output = ResNet(b_x)  # 先将数据放到cnn中计算output
        loss = loss_func(output, b_y)  # 输出和真实标签的loss，二者位置不可颠倒

        total_train_loss += loss.data.item()
        optimizer.zero_grad()  # 清除之前学到的梯度的参数
        loss.backward()  # 反向传播，计算梯度
        optimizer.step()  # 应用梯度
        _, train_pred_y = torch.max(output, 1)
        train_accuracy = (train_pred_y == b_y).sum()
        total_train_acc += train_accuracy.item()

    print("整体训练集上的loss：{}，正确率：{:.2f}%".format(total_train_loss / train_size, 100 * total_train_acc / train_size))
    train_loss.append(total_train_loss / train_size)
    train_acc.append(100 * total_train_acc / train_size)

    for test_step, (test_x, test_y) in enumerate(test_dataloader):
        if cuda_is_available:
            test_x = test_x.cuda()
            test_y = test_y.cuda()
        test_output = ResNet(test_x)
        _, pred_y = torch.max(test_output, 1)
        test_loss_1 = loss_func(test_output, test_y)
        total_test_loss += test_loss_1.data.item()
        accuracy = (pred_y == test_y).sum()
        total_test_acc += accuracy.item()
        # print('| test loss: %.4f' % test_loss.data.numpy(), '| test accuracy: %.2f' % accuracy)
    print("整体测试集上的loss：{}，正确率：{:.2f}%".format(total_test_loss / test_size, 100 * total_test_acc / test_size))
    test_loss.append(total_test_loss / test_size)
    test_acc.append(100 * total_test_acc / test_size)
# torch.save(cnn.state_dict(), 'cnn2.pkl')#保存模型

# 加载模型，调用时需将前面训练及保存模型的代码注释掉，否则会再训练一遍
# cnn.load_state_dict(torch.load('cnn2.pkl'))
# cnn.eval()


x = range(EPOCHS)
ax = plt.gca()
plt.plot(x, train_loss, 'b', label="Train_loss")
plt.plot(x, test_loss, 'r', label="test_loss")
plt.title("Train and Test loss")
plt.legend(loc='upper right')
plt.figure()

plt.plot(x, train_acc, 'b', label="Train_accuracy")
plt.plot(x, test_acc, 'r', label="Test_accuracy")
plt.title("Train and Test accuracy")
plt.legend(loc="lower right")
plt.show()