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 SummaryWriterimport 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 = Falseself.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 xnet = 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.0print('-' * 10)running_loss = 0.0running_corrects = 0model = 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.0train_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.0running_corrects = 0model = 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的lossrunning_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)