Resnet训练 验证自己的数据集

使用resnet神经网络，在预训练模型的基础上，在自己的数据集上进行微调，最后在自己的数据集上验证效果。

可以借鉴部分：

• 数据加载
• resnet模型使用方法
• 预训练模型加载，模型保存
• 训练/验证/测试步骤
• 使用cpu/gpu进行模型训练

需要修改部分：

• 预训练模型加载路径
• 自己数据集路径
• 保存路径

import torch.nn as nn
import math
import pickle
import torch
import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import xlwt
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
import numpy as np
from torchvision import models
import matplotlib.pyplot as plt
from PIL import Image
__all__ = ['ResNet', 'resnet50']


def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)

class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class ResNet(nn.Module):

    def __init__(self, block, layers, num_classes=15, include_top=True):
        self.inplanes = 64
        super(ResNet, self).__init__()
        self.include_top = include_top

        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=0, ceil_mode=True)

        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7, stride=1)
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)

        if not self.include_top:
            return x

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


def resnet50(**kwargs):
    """Constructs a ResNet-50 model.
    """
    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
    return model

def load_state_dict(model, fname):
    """
    Set parameters converted from Caffe models authors of VGGFace2 provide.
    See https://www.robots.ox.ac.uk/~vgg/data/vgg_face2/.

    Arguments:
        model: model
        fname: file name of parameters converted from a Caffe model, assuming the file format is Pickle.
    """
    with open(fname, 'rb') as f:
        weights = pickle.load(f, encoding='latin1')

    own_state = model.state_dict()
    #print(own_state)
    for name, param in weights.items():
        #print(name)
        if name in own_state:
            try:
                own_state[name].copy_(torch.from_numpy(param))
            except Exception:
                raise RuntimeError('While copying the parameter named {}, whose dimensions in the model are {} and whose '\
                                   'dimensions in the checkpoint are {}.'.format(name, own_state[name].size(), param.size()))
        else:
            raise KeyError('unexpected key "{}" in state_dict'.format(name))

if __name__=='__main__':
    #model = torchvision.models.resnet50()

    weight_file='/train/resNet/resnet50_scratch_weight.pkl'
    # weight_file = './resnet50_scratch_weight.pkl'
    model_ft = models.resnet50()
    # print(model_ft)
    model_ft.avgpool = nn.AvgPool2d(kernel_size=7, stride=1, padding=0)
    model_ft.fc = nn.Linear(2048, 8631)
    load_state_dict(model_ft, weight_file)
    model_ft.fc = nn.Linear(2048, 15)
    num=0
    # print('---'*10)
    # print(model_ft)
    # for i in model_ft.parameters():
    #     num=num+1
    #
    #     print(i)
    batch_size=16
    train_transforms = transforms.Compose([
        transforms.Resize((224,224)),  # 对图片尺寸做一个缩放切割
        transforms.RandomHorizontalFlip(),  # 水平翻转
        transforms.Grayscale(num_output_channels=3),
        transforms.ToTensor(),  # 转化为张量
        # transforms.Normalize((.5, .5, .5), (.5, .5, .5))  # 进行归一化
    ])
    # 对测试集做变换
    val_transforms = transforms.Compose([
        transforms.Resize((224,224)),
        # transforms.RandomResizedCrop(224),
        transforms.Grayscale(num_output_channels=3),
        transforms.ToTensor(),
        # transforms.Normalize((.5, .5, .5), (.5, .5, .5))
    ])

    train_dir = "/train"  # 训练集路径
    train_datasets = datasets.ImageFolder(train_dir, transform=train_transforms)
    # 加载数据集
    train_dataloader = torch.utils.data.DataLoader(train_datasets, batch_size=batch_size, shuffle=True)

    val_dir = "/val"
    val_datasets = datasets.ImageFolder(val_dir, transform=val_transforms)
    val_dataloader = torch.utils.data.DataLoader(val_datasets, batch_size=batch_size, shuffle=True)

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(train_datasets.class_to_idx)

    # Hyper-parameters
    num_epochs = 800
    learning_rate = 0.001
    if torch.cuda.is_available():
        model_ft.cuda()
    #     print('++++++++++++++++')
    max_acc=0
    # params = [{'params': md.parameters()} for md in model_ft.children()
    #           if md in [model_ft.classifier]]
    optimizer = optim.SGD(filter(lambda p: p.requires_grad, model_ft.parameters()), lr=learning_rate)
    loss_func = nn.CrossEntropyLoss()

    for epoch in range(200):  # 100
        print('epoch= ',epoch)
        # training-----------------------------
        model_ft.train()
        train_loss = 0.
        train_acc = 0.
        if(epoch!=0 and epoch%10==0):
            learning_rate=learning_rate*0.1
            optimizer = optim.SGD(filter(lambda p: p.requires_grad, model_ft.parameters()), lr=learning_rate)

        for batch_x, batch_y in train_dataloader:
            batch_x, batch_y = Variable(batch_x).cuda(), Variable(batch_y).cuda()
            out = model_ft(batch_x)
            loss = loss_func(out, batch_y)
            train_loss += loss.item()
            pred = torch.max(out, 1)[1]
            # print(pred)
            train_correct = (pred == batch_y).sum()
            train_acc += train_correct.item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        strr=str(epoch)+"  "+str(train_loss/len(train_datasets)*100)+'    '+str(train_acc/len(train_datasets)*100)
        with open('s_loss.txt','a') as f:
            f.write(strr+'\n')
        # evaluation--------------------------------
        # if(epoch%20==0):
        if(True):
            model_ft.eval()
            eval_loss = 0.
            eval_acc = 0.
            for batch_x, batch_y in val_dataloader:
                batch_x, batch_y = Variable(batch_x, volatile=True).cuda(), Variable(batch_y, volatile=True).cuda()
                out =model_ft(batch_x)
                loss = loss_func(out, batch_y)
                eval_loss += loss.item()
                pred = torch.max(out, 1)[1]
                # print('-------------')
                # print(pred)
                num_correct = (pred == batch_y).sum()
                eval_acc += num_correct.item()
            leng = len(val_datasets)
            leng=1.0*leng
            tmpp=eval_acc/leng*100
            print('eval_acc ',tmpp)
            strr=str(eval_acc)+"  "+str(tmpp)
            with open('s_eval_acc.txt','a') as f:
                f.write(strr+'\n')
            if(tmpp>max_acc):
                max_acc=tmpp
                torch.save(model_ft.state_dict(), '/train/resNet/resnet_s_best.pkl')
    test_dir = "/test"
    test_datasets = datasets.ImageFolder(test_dir, transform=val_transforms)
    test_dataloader = torch.utils.data.DataLoader(test_datasets, batch_size=batch_size, shuffle=True)
    test_loss = 0.
    test_acc = 0.
    for batch_x, batch_y in test_dataloader:
        batch_x, batch_y = Variable(batch_x, volatile=True).cuda(), Variable(batch_y, volatile=True).cuda()
        out = model_ft(batch_x)
        pred = torch.max(out, 1)[1]
        num_correct = (pred == batch_y).sum()
        test_acc += num_correct.item()
        str1 = pred.cuda().data.cpu().numpy()
        str2 = batch_y.cuda().data.cpu().numpy()
        with open('an_est2.txt', 'a') as f:
            f.write(str(str1) + '\n')
            f.write(str(str2) + '\n')
    leng = len(test_datasets)
    leng = 1.0 * leng
    tmpp = test_acc / leng * 100
    print('eval_acc2 ', tmpp, test_acc, leng)
    strr = str(test_acc) + "  " + str(tmpp)
    with open('s_eval_test2.txt', 'a') as f:
        f.write(strr + '\n')