利用AlexNet网络进行训练-没有废话-全是干货-完整代码

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

文章目录

• AlexNet网络
• 训练代码

---

前言

AlexNet网络进行图片分类识别。

---

import torch
import torchvision
#import torchvision.models

from torchvision.models import AlexNet
from matplotlib import pyplot as plt
from tqdm import tqdm
from torch import nn
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
from torch.utils.tensorboard import SummaryWriter
from PIL import ImageFile
tensorboardwriter = SummaryWriter("./logs")
writer = SummaryWriter("./logs")
'''x = torch.rand(8, 3, 256, 512)
model = AlexNet()'''

'''with SummaryWriter(comment='AlexNet') as w:
    w.add_graph(model, x)'''
data_transform = {
    "train": transforms.Compose([transforms.RandomResizedCrop(100),
                                 transforms.RandomHorizontalFlip(),
                                 transforms.ToTensor(),
                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
    "val": transforms.Compose([transforms.Resize((100, 100)),  # cannot 224, must (224, 224)
                               transforms.ToTensor(),
                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}




def main():
    train_data = torchvision.datasets.ImageFolder(root=r"train", transform=data_transform["train"])

    traindata = DataLoader(dataset=train_data, batch_size=128, shuffle=True, num_workers=0)

    test_data = torchvision.datasets.ImageFolder(root=r"val", transform=data_transform["val"])

    train_size = len(train_data)
    test_size = len(test_data)
    print(train_size)
    print(test_size)
    testdata = DataLoader(dataset=test_data, batch_size=128, shuffle=True, num_workers=0)

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    class alexnet(nn.Module):
        def __init__(self):
            super(alexnet, self).__init__()
            self.model = nn.Sequential(

                nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2),  # input[3, 120, 120]  output[48, 55, 55]

                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=3, stride=2),  # output[48, 27, 27]
                nn.Conv2d(48, 128, kernel_size=5, padding=2),  # output[128, 27, 27]

                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=3, stride=2),  # output[128, 13, 13]
                nn.Conv2d(128, 192, kernel_size=3, padding=1),  # output[192, 13, 13]

                nn.ReLU(inplace=True),
                nn.Conv2d(192, 192, kernel_size=3, padding=1),  # output[192, 13, 13]

                nn.ReLU(inplace=True),
                nn.Conv2d(192, 128, kernel_size=3, padding=1),  # output[128, 13, 13]

                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=3, stride=2),  # output[128, 6, 6]
                nn.Flatten(),
                nn.Dropout(p=0.5),
                nn.Linear(512, 2048),
                nn.ReLU(inplace=True),
                nn.Dropout(p=0.5),
                nn.Linear(2048, 1024),
                nn.ReLU(inplace=True),
                nn.Linear(1024, 5),  # 自己的数据是几种，这里的7就改成几

            )

        def forward(self, x):
            x = self.model(x)
            return x

    alexnet1 = alexnet()
    '''x = torch.rand(8, 3, 256, 512)
    #model = AlexNet()
    with SummaryWriter(comment='alexnet') as w:
        w.add_graph(alexnet1, x)'''
    print(alexnet1)
    alexnet1.to(device)
    test1 = torch.ones(64, 3, 100, 100)

    test1 = alexnet1(test1.to(device))
    print(test1.shape)

    epoch = 50
    learning = 0.0001
    optimizer = torch.optim.Adam(alexnet1.parameters(), lr=learning)
    loss = nn.CrossEntropyLoss()

    train_loss_all = []
    train_accur_all = []
    test_loss_all = []
    test_accur_all = []
    for i in range(epoch):
        train_loss = 0
        train_num = 0.0
        train_accuracy = 0.0
        alexnet1.train()
        train_bar = tqdm(traindata)
        for step, data in enumerate(train_bar):
            img, target = data
            optimizer.zero_grad()
            outputs = alexnet1(img.to(device))

            loss1 = loss(outputs, target.to(device))
            outputs = torch.argmax(outputs, 1)
            loss1.backward()
            optimizer.step()
            train_loss += abs(loss1.item()) * img.size(0)
            accuracy = torch.sum(outputs == target.to(device))
            train_accuracy = train_accuracy + accuracy
            train_num += img.size(0)

        print("epoch：{} ， train-Loss：{} , train-accuracy：{}".format(i + 1, train_loss / train_num,
                                                                    train_accuracy / train_num))
        train_loss_all.append(train_loss / train_num)
        train_accur_all.append(train_accuracy.double().item() / train_num)
        test_loss = 0
        test_accuracy = 0.0
        test_num = 0
        alexnet1.eval()
        with torch.no_grad():
            test_bar = tqdm(testdata)
            for data in test_bar:
                img, target = data

                outputs = alexnet1(img.to(device))

                loss2 = loss(outputs, target.to(device))
                outputs = torch.argmax(outputs, 1)
                test_loss = test_loss + abs(loss2.item()) * img.size(0)
                accuracy = torch.sum(outputs == target.to(device))
                test_accuracy = test_accuracy + accuracy
                test_num += img.size(0)

        print("test-Loss：{} , test-accuracy：{}".format(test_loss / test_num, test_accuracy / test_num))
        test_loss_all.append(test_loss / test_num)
        test_accur_all.append(test_accuracy.double().item() / test_num)
    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)

    plt.plot(range(epoch), train_loss_all,
             label="Train loss")
    plt.plot(range(epoch), test_loss_all,
             label="test loss")
    plt.legend()
    plt.xlabel("epoch")
    plt.ylabel("Loss")
    plt.subplot(1, 2, 2)
    plt.plot(range(epoch), train_accur_all,
             label="Train accur")
    plt.plot(range(epoch), test_accur_all,
             label="test accur")
    plt.xlabel("epoch")
    plt.ylabel("acc")

    plt.legend(['train', 'test'], loc='upper left')
    plt.legend()
    plt.savefig('train.png')
    torch.save(alexnet1.state_dict(), "alexnet.pth")

    print("模型已保存")


if __name__ == '__main__':
    main()

例如：以上就是今天要讲的内容，本文仅仅简单介绍了pandas的使用，而pandas提供了大量能使我们快速便捷地处理数据的函数和方法。