#!/usr/bin/env python # encoding: utf-8 ''' @author: liualex @contact: [email protected] @software: pycharm @file: main.py @time: 2019/8/16 16:21 @desc: ''' import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms import torch.utils.data.dataloader as dataloader import pdb import os os.environ["CUDA_VISIBLE_DEVICES"] = "3" transform = transforms.Compose( [ transforms.RandomHorizontalFlip(), transforms.RandomGrayscale(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ] ) train_set = torchvision.datasets.CIFAR10( root="./data", download=True, train=True, transform=transform ) train_loader = dataloader.DataLoader( dataset=train_set, batch_size=100, shuffle=False ) test_set = torchvision.datasets.CIFAR10( root="./data", download=True, train=False, transform=transform ) test_loader = dataloader.DataLoader( dataset=test_set, batch_size=100, shuffle=True ) class AlexNet(nn.Module): def __init__(self, class_num): super(AlexNet, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 192, kernel_size=3, stride=1, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(192, 384, kernel_size=3, stride=1, padding=1), nn.ReLU(inplace=True), nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1), nn.ReLU(inplace=True), nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), ) self.classifier = nn.Sequential( nn.Dropout(), nn.Linear(256 * 4 * 4, 4096), nn.ReLU(inplace=True), nn.Dropout(), nn.Linear(4096, 4096), nn.ReLU(inplace=True), nn.Linear(4096, class_num), ) def forward(self, x): x = self.features(x) x = x.view(x.size(0), 256 * 4 * 4) x = self.classifier(x) return x device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = AlexNet(10) model.to(device) epoches = 100 lr = 0.001 criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=lr) for epoch in range(epoches): for i, (images, labels) in enumerate(train_loader): images = images.to(device) labels = labels.to(device) output = model(images) loss = criterion(output, labels) optimizer.zero_grad() loss.backward() optimizer.step() if (i + 1) % 100 == 0: print('Epoch [{}/{}], Loss: {:.4f}' .format(epoch + 1, epoches, loss.item())) with torch.no_grad(): correct = 0 total = 0 for (images, labels) in test_loader: images, labels = images.to(device), labels.to(device) output = model(images) _, predicted = torch.max(output.data, 1) total += labels.size(0) #pdb.set_trace() correct += (predicted == labels).sum().item() print("The accuracy of total {} images: {}%".format(total, 100 * correct / total))
Convolution calculation formula: N = ( W - F. + 2 P ) / S + 1'd , wherein W refers to the size of the input image, F refers to the size of the convolution kernel, P refers to the padding, S refers to stride
Results: Since only iterated 100 times, the results of only 78.57%.