pytorch学习之路——07

建立cnn模型

import torch
from torch import nn,optim
import torch.nn.functional as F
from torchvision import datasets,transforms
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler

import numpy as np



def get_dataloader(batch_size,valid_size):
    transform = transforms.Compose([transforms.ToTensor(),
                                    transforms.Normalize([0.5,],
                                                         [0.5,])])
    trainset = datasets.MNIST('MNIST_data/',train=True,transform=transform)
    testset = datasets.MNIST('MNIST_data/',train=False,transform=transform)

    num_train = len(trainset)
    indexs = list(range(num_train))
    np.random.shuffle(indexs)
    split = int(np.floor(valid_size*num_train))
    train_idx,valid_idx = indexs[split:],indexs[:split]

    train_sampler = SubsetRandomSampler(train_idx)
    valid_sampler = SubsetRandomSampler(valid_idx)
    train_loader = DataLoader(trainset,batch_size=batch_size,sampler=train_sampler)
    valid_loader = DataLoader(trainset,batch_size=batch_size,sampler=valid_sampler)
    test_loader = DataLoader(testset,batch_size=batch_size)

    return train_loader,valid_loader,test_loader


class Network(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1,8,3,padding=1)
        self.conv2 = nn.Conv2d(8,16,3,padding=1)
        self.conv3 = nn.Conv2d(16,32,3,padding=1)
        self.pool = nn.MaxPool2d(2,2)
        self.fc1 = nn.Linear(32*3*3,64)
        self.fc2 = nn.Linear(64,10)
        self.drop = nn.Dropout(p=0.3)

    def forward(self,x):
        x = self.pool(self.conv1(x))
        x = self.pool(self.conv2(x))
        x = self.pool(self.conv3(x))
        x = x.view(-1,32*3*3)
        x = F.relu(self.fc1(x))
        x = self.drop(x)
        x = F.log_softmax(self.fc2(x),dim=1)

        return x


def train(train_loader,valid_loader,epoch,device):
    model = Network()
    model.to(device)
    costFunc = nn.NLLLoss()
    optimizer = optim.SGD(model.parameters(),lr=0.03)

    min_valid_loss = np.Inf
    for e in range(epoch):
        train_loss = 0
        for imgs,labs in train_loader:
            imgs,labs = imgs.to(device),labs.to(device)
            optimizer.zero_grad()
            logps = model.forward(imgs)
            loss = costFunc(logps,labs)
            loss.backward()
            optimizer.step()
            train_loss += loss

        valid_loss = 0
        with torch.no_grad():
            for imgs,labs in valid_loader:
                imgs,labs = imgs.to(device),labs.to(device)
                logps = model.forward(imgs)
                loss = costFunc(logps,labs)
                valid_loss += loss

        train_loss /= len(train_loader)
        valid_loss /= len(valid_loader)
        print("Epoch :{}/{},".format(e,epoch),
              "Train Loss:{:.3f}".format(train_loss),
              "Valid Loss:{:.3f}".format(valid_loss))
        if valid_loss <= min_valid_loss:
            print("Epoch: {}".format(e),
                  "Valid Loss Decrease {:.3f} -> {:.3f}".format(min_valid_loss,valid_loss))
            min_valid_loss = valid_loss
            torch.save(model.state_dict(),"minst_cnn.pth")
            best_model = model

    return best_model


def predict(path,test_loader,device):
    model = Network()
    costFunc = nn.NLLLoss()
    state_dict = torch.load(path)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()

    test_loss = 0
    class_correct = np.zeros(10)
    class_total = np.zeros(10)
    accuracy = 0
    for imgs,labs in test_loader:
        imgs,labs = imgs.to(device),labs.to(device)
        logps = model.forward(imgs)
        loss = costFunc(logps,labs)
        test_loss += loss

        ps = torch.exp(logps)
        top_k,top_class = ps.topk(1,dim=1)
        equal_tensor = top_class == labs.view(*top_class.shape)
        equal = np.squeeze(equal_tensor.numpy()) if device is not "gpu" else np.squeeze(equal_tensor.cpu().numpy())
        accuracy += np.mean(equal)
        for i in range(min(batch_size,labs.shape[0])):
            lab = labs[i]
            class_correct[lab] += equal[i]
            class_total[lab] += 1

    for i in range(10):
        if class_total[i] > 0:
            print("Accuracy:{:.3f} ({}/{})".format(class_correct[i]/class_total[i],class_correct[i],class_total[i]))

    print("Test Accuracy (Overall):",np.sum(class_correct)/np.sum(class_total))
    print(accuracy/len(test_loader))
    print(test_loss/len(test_loader))


if __name__ == "__main__":
    batch_size = 64
    valid_size = 0.2
    epoch = 5
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_loader,valid_loader,test_loader = get_dataloader(batch_size,valid_size)
    # model = train(train_loader,valid_loader,epoch,device)
    predict('minst_cnn.pth',test_loader,device)