Pytorch图像分类任务模板(cifar10-94%)

train.py:

# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import numpy as np
from torchvision import models
import torchvision.transforms as transforms
from cifar10 import CIFAR10

# define network
def genNet():
    net=models.resnet34(pretrained=True) # pretrained
    net.conv1=nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) #(b,64,32,32)
    net.bn1=nn.BatchNorm2d(64)
    net.maxpool=nn.MaxPool2d(kernel_size=1, stride=1)
    net.avgpool=nn.AvgPool2d(4, stride=1) #32/8
    expansion=1 #18,34 : expansion=1 else 4
    net.fc=nn.Linear(512 * expansion, 10) #
    return net

# define parameters
torch.cuda.set_device(0)
gpu_en=True
load=True
save=True
num_epochs=1
nbatch=150
lr=0.001

# define transform
transform=transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# define dataset
dataset=CIFAR10('D:\\dataset\\cifar10', train=True, transform=transform)

# run
net=genNet()
if load:
    net.load_state_dict(torch.load('net.pkl'))
if gpu_en:
    net=net.cuda()
loader=torch.utils.data.DataLoader(dataset, batch_size=nbatch, shuffle=True, num_workers=0)
criterion=nn.CrossEntropyLoss()
optimizer=torch.optim.SGD(net.parameters(), lr=lr)
net.train()
for epoch in range(num_epochs):
    for i, (inputs, labels) in enumerate(loader):
        if gpu_en:
            inputs=inputs.cuda()
            labels=labels.cuda()
        optimizer.zero_grad()
        outputs=net(inputs)
        loss=criterion(outputs, labels)
        Y_pred=torch.max(outputs,1)[1]
        acc=float(sum(np.array(Y_pred==labels)))/len(labels)
        print('epc:%d,step:%d,loss:%f,acc:%f' % (epoch,i,loss,acc))
        loss.backward()
        optimizer.step()
if save:
    torch.save(net.state_dict(),'net.pkl')

test.py:

# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import numpy as np
from torchvision import models
import torchvision.transforms as transforms
from cifar10 import CIFAR10

# define network
def genNet():
    net=models.resnet34(pretrained=False)
    net.conv1=nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) #(b,64,32,32)
    net.bn1=nn.BatchNorm2d(64)
    net.maxpool=nn.MaxPool2d(kernel_size=1, stride=1)
    net.avgpool=nn.AvgPool2d(4, stride=1) #32/8
    expansion=1 #18,34 : expansion=1 else 4
    net.fc=nn.Linear(512*expansion, 10) #
    return net

# define parameters
torch.cuda.set_device(0)
gpu_en=True
nbatch=200

# define transform
transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# define dataset
dataset=CIFAR10('D:\\dataset\\cifar10', train=False, transform=transform)

# run
net=genNet()
net.load_state_dict(torch.load('net.pkl'))
if gpu_en:
    net=net.cuda()
loader=torch.utils.data.DataLoader(dataset, batch_size=nbatch, shuffle=False, num_workers=0)
net.eval()
for epoch in range(1):
    acc_sum=0.0
    idx=0
    for i, (inputs, labels) in enumerate(loader):
        if gpu_en:
            inputs=inputs.cuda()
            labels=labels.cuda()
        outputs=net(inputs)
        Y_pred=torch.max(outputs,1)[1]
        acc=float(sum(np.array(Y_pred==labels)))/len(labels)
        acc_sum+=acc
        idx+=1
        print('step:%d, acc:%f' % (i, acc))
    print('acc_avg:%f' % (acc_sum/idx))

cifar10.py:

# -*- coding: utf-8 -*-
import os
import numpy as np
import matplotlib.pyplot as plt
import torch.utils.data as data
from PIL import Image
import random
import torchvision.transforms as transforms
import torch

# require format: binary

class CIFAR10(data.Dataset):

    def __init__(self, root, train=True, transform=None, target_transform=None):

        self.dic = ['airplane','automobile','bird','cat','deer','dog','frog','horse','ship' ,'truck']
        self.root = root
        self.transform = transform
        self.target_transform = target_transform
        self.train = train  # training set or test set

        # imgs:np.uint8, labels:np.int64
        if self.train:

            f1 = open(os.path.join(root, 'data_batch_1.bin'), 'rb')
            f2 = open(os.path.join(root, 'data_batch_2.bin'), 'rb')
            f3 = open(os.path.join(root, 'data_batch_3.bin'), 'rb')
            f4 = open(os.path.join(root, 'data_batch_4.bin'), 'rb')
            f5 = open(os.path.join(root, 'data_batch_5.bin'), 'rb')
            raw = f1.read()+f2.read()+f3.read()+f4.read()+f5.read()
            self.labels = []
            self.data = []

            for i in range(50000):
                nidx = 3073*i
                labels_np=np.array(list(raw[nidx:nidx+1]),dtype='int64')
                data_np=np.array(list(raw[nidx+1:nidx+3073]),dtype='uint8')
                self.labels.append(labels_np)
                self.data.append(data_np)
            self.labels = np.concatenate(self.labels)
            self.data = np.concatenate(self.data)

            self.data = self.data.reshape((50000, 3, 32, 32))
            self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC
            f1.close()
            f2.close()
            f3.close()
            f4.close()
            f5.close()

        else:

            f1 = open(os.path.join(root, 'test_batch.bin'), 'rb')
            raw = f1.read()
            self.labels = []
            self.data = []

            for i in range(10000):
                nidx = 3073*i
                labels_np=np.array(list(raw[nidx:nidx+1]),dtype='int64')
                data_np=np.array(list(raw[nidx+1:nidx+3073]),dtype='uint8')
                self.labels.append(labels_np)
                self.data.append(data_np)
            self.labels = np.concatenate(self.labels)
            self.data = np.concatenate(self.data)

            self.data = self.data.reshape((10000, 3, 32, 32))
            self.data = self.data.transpose((0, 2, 3, 1))  # convert to HWC
            f1.close()

    def __getitem__(self, index):

        if self.train:
            img, target = self.data[index], self.labels[index]
        else:
            img, target = self.data[index], self.labels[index]

        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self):

        if self.train:
            return 50000
        else:
            return 10000

    def test(self):
        if self.train:
            for i in range(15):
                plt.subplot(3,5,i+1)
                idx=random.randint(0,50000-1)
                img,label=self.__getitem__(idx)
                plt.axis('off')
                plt.imshow(img)
                plt.title(self.dic[label])
        else:
            for i in range(15):
                plt.subplot(3,5,i+1)
                idx=random.randint(0,10000-1)
                img,label=self.__getitem__(idx)
                plt.axis('off')
                plt.imshow(img)
                plt.title(self.dic[label])
        plt.show()


if __name__ == '__main__':

    trans_en=True

    if trans_en:
        transform = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
        ])
    else:
        transform=None

    DS=CIFAR10('D:\\dataset\\cifar10',train=False,transform=transform)
    DS.test()

转载于:https://www.cnblogs.com/xytpai/p/9614745.html