CNN的参数计算

CNN图像分类

import torch
import torch.nn as nn
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor
import torch.optim as optim
from torch.utils.data import DataLoader
import time
import matplotlib.pyplot as plt
from torchsummary import  summary

BATCH_SIZE = 8
#1.准备数据集
def create_dataset():

    train_dataset = CIFAR10('./data',train=True,transform=ToTensor(),download=True)
    test_dataset = CIFAR10('./data',train=False,transform=ToTensor(),download=True)

    return train_dataset,test_dataset
#2。搭建神经网络
class ImgModel(nn.Module):

    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(3,6,3,1,0)
        self.pool1 = nn.MaxPool2d(2,2,0)
        self.conv2 = nn.Conv2d(6,6,3,1,0)
        self.pool2 = nn.MaxPool2d(2,2,0)

        self.linear1 = nn.Linear(216,120)
        self.linear2 = nn.Linear(120,84)
        self.output = nn.Linear(84,10)

    def forward(self,x):
        #第一层：卷积层（加权求和）-》激励层（激活函数）+池化层
        x = self.pool1(torch.relu(self.conv1(x)))
        #第二层
        x = self.pool2(torch.relu(self.conv2(x)))
        #全连接层：细节——>全连接层只能处理2维数据
        x = x.reshape(x.size(0),-1)

        x =torch.relu(self.linear1(x))

        x =torch.relu(self.linear2(x))

        #输出层
        return self.output(x) #不用softmax，因为后续的损失函数用CrossEntropyLoss
def train(train_dataset):
    dataloader = DataLoader(train_dataset,batch_size=BATCH_SIZE,shuffle=True)
    model = ImgModel()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(),lr=0.001)
    epochs = 10
    for epoch in range(epochs):
        total_loss,total_samples,total_correct,start = 0.0,0,0,time.time()
        for x,y in dataloader:
            model.train()
            y_pred = model(x)
            loss=criterion(y_pred,y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_correct+=(torch.argmax(y_pred,dim=-1)==y).sum()
            total_loss+=loss.item()*len(y)
            total_samples +=  len(y)
        print(f'epoch:{epoch+1},loss:{total_loss/total_samples:.4f},accuracy:{total_correct/total_samples:.2f},time:{time.time()-start:.2f}s')
    torch.save(model.state_dict(),'./model/img_model.pth')
def evaluate():
    dataloader = DataLoader(test_dataset,batch_size=BATCH_SIZE,shuffle=False)

    model = ImgModel()

    model.load_state_dict(torch.load('./model/img_model.pth'))

    total_correct,total_samples = 0,0

    for x,y in dataloader:
        model.eval()

        y_pred = model(x)

        #搭建网络时，没有用softmax函数，用argmax函数

        y_pred = torch.argmax(y_pred,dim=-1)
        total_correct+=(y_pred == y).sum()

        total_samples+=len(y)

    print(f'accuracy:{total_correct/total_samples:.2f}')

if __name__ == '__main__':
    #1.获取数据集
    train_dataset,test_dataset = create_dataset()
    #2.搭建神经网络
    model = ImgModel()
    #查看模型参数
    # summary(model,(3,32,32),batch_size=BATCH_SIZE)
    #训练
    train(train_dataset)
    evaluate()