pytorch实现图像分类代码实例

图像多标签分类例子

import os
import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator
from tensorboardX import SummaryWriter

import seaborn as sns
from sklearn.metrics import confusion_matrix

'''数据加载'''
#选择设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#对三种数据集进行不同预处理，对训练数据进行加强
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomRotation(30),
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ]),
    'valid': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
}
 
#数据目录
data_dir = "/DATA/wanghongzhi/17flowers"
 
#获取两个数据集
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),  #要习惯python这种语法
                 data_transforms[x]) for x in ['train', 'valid']} 
traindataset = image_datasets['train']
validdataset = image_datasets['valid']
 
batch_size = 8
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size,
             shuffle=True, num_workers=4) for x in ['train', 'valid']}
print(dataloaders)
traindataloader = dataloaders['train']
validdataloader = dataloaders['valid']
 
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid']}

'''定义网络结构'''
class Net(nn.Module):
    def __init__(self,model):
        super(Net,self).__init__()
        self.features = model.features
        # for p in self.parameters():
        #     p.requires_grad = False
        self.classifier = nn.Sequential(
            nn.Linear(25088, 4096,bias=True),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5,inplace=False),
            nn.Linear(4096, 4096,bias=True),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5,inplace=False),
            nn.Linear(4096, 102,bias=True)
        )
 
    def forward(self,x):
        x = self.features(x)
        x = x.view(x.shape[0], -1)
        x = self.classifier(x)
        return x
 
net = models.resnet50().to(device)

net.load_state_dict(torch.load('/home/wanghongzhi/zuoye/resnet50.pth'))

'''参数设定'''
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, net.parameters()),lr=0.0001,momentum=0.9)

'''定义根据loss列表绘制loss曲线函数'''
def hua_loss(loss):
    l=len(loss)
    #x=list(range(1,l+1))
    x=range(1,l+1)

    # 设置图片大小
    plt.figure(figsize=(20,8),dpi=80) # figsize设置图片大小，dpi设置清晰度
    plt.title("Train-Epoch-Loss",fontsize=25)
    plt.xlabel("Epoch",fontsize=20)
    plt.ylabel("Loss",fontsize=20)
    plt.plot(x,loss)
    x_major_locator=MultipleLocator(2)  #x轴刻度为1的倍数
    y_major_locator=MultipleLocator(0.15) #y轴刻度为0.01的倍数
    ax=plt.gca() #ax为两条坐标轴的实例
    ax.xaxis.set_major_locator(x_major_locator)
    ax.yaxis.set_major_locator(y_major_locator)
    #保存
    #plt.savefig("./t1.png")
    plt.show()

'''先定义验证集检验''' #测试集和验证集代码一模一样
def valid_model(model, criterion):
    best_acc = 0.0
    print('-' * 10)
 
    running_loss = 0.0
    running_corrects = 0
    model = model.to(device)
    for inputs, labels in validdataloader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        model.eval()
        with torch.no_grad():
            outputs = model(inputs)
        loss = criterion(outputs, labels)
        print('outputs:',outputs)
        print('labels:',labels)
        _, preds = torch.max(outputs, 1)
        running_loss += loss.item()
        running_corrects += torch.sum(preds == labels).item()
    epoch_loss = running_loss / dataset_sizes['valid']
    print(running_corrects)
    epoch_acc = running_corrects / dataset_sizes['valid']
    print('{} Loss: {:.4f} Acc: {:.4f}'.format(
            'valid', epoch_loss, epoch_acc))
    print('-' * 10)
    print()
    #val_loss.append(epoch_loss)

'''训练模型'''
def train_model(model, criterion, optimizer, num_epochs=5):
    #since = time.time()
    best_acc = 0.0
    train_loss=[]
    #val_loss=[]

    for epoch in range(num_epochs):
        if (epoch+1)%5==0:       #每五个epoch就用该模型验证一次结果          
            valid_model(model, criterion)
        print('-' * 10)
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
 
        running_loss = 0.0
        running_corrects = 0
        model = model.to(device)
        for inputs, labels in traindataloader:
            inputs = inputs.to(device)
            labels = labels.to(device)
            model.train()
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
 
            _, preds = torch.max(outputs, 1)
            running_loss += loss.item()  #加起来用来计算每个epoch的loss
            running_corrects += torch.sum(preds == labels).item() #item()取出张量中的值，或者(predicted==labels).sum().item()
        epoch_loss = running_loss / dataset_sizes['train']
        print(dataset_sizes['train'])  #训练集总数
        print(running_corrects) #正确预测个数
        epoch_acc = running_corrects / dataset_sizes['train']
        best_acc = max(best_acc,epoch_acc)
        print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                'train', epoch_loss, epoch_acc)) 
        print()
        train_loss.append(epoch_loss)

    hua_loss(train_loss)

 
    print('Best val Acc: {:4f}'.format(best_acc)) 
    return model



'''开始训练'''
epochs = 5
model = train_model(net, criterion, optimizer, epochs)

参考了博客：Pytorch实现鲜花分类(102 Category Flower Dataset)