|
- import torch
- import torchvision
- from torch.utils.data import DataLoader
- import torch.nn as nn
- import torch.optim as optim
-
- lr=0.001
- transform_train = torchvision.transforms.Compose([torchvision.transforms.Pad(4),
- torchvision.transforms.RandomHorizontalFlip(),
- torchvision.transforms.RandomCrop(32),
- torchvision.transforms.ToTensor(),
- torchvision.transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
-
- transform_test = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),
- torchvision.transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
-
- trainSet=torchvision.datasets.CIFAR10("/dataset",train=True,download=False,transform=transform_train)
- testSet=torchvision.datasets.CIFAR10("/dataset",train=False,download=False,transform=transform_test)
- trainLoader=DataLoader(trainSet,batch_size=32,shuffle=True,num_workers=4)
- testLoader=DataLoader(testSet,batch_size=32,shuffle=False,num_workers=4)
-
-
- device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
- class CNNNet(nn.Module):
- def __init__(self):
- super(CNNNet, self).__init__()
- self.conv1=nn.Sequential(nn.Conv2d(3,16,5),nn.MaxPool2d(2,2),nn.ReLU())
- self.conv2=nn.Sequential(nn.Conv2d(16,36,3),nn.MaxPool2d(2,2),nn.ReLU())
- self.fc1=nn.Sequential(nn.Linear(1296,128),nn.ReLU(),nn.BatchNorm1d(128))
- self.fc2=nn.Sequential(nn.Linear(128,84),nn.ReLU(),nn.BatchNorm1d(84))
- self.fc3 = nn.Sequential(nn.Linear(84, 10))
-
- def forward(self,x):
- x=self.conv1(x)
- x=self.conv2(x)
- x=x.view(-1,1296)
- x=self.fc1(x)
- x=self.fc2(x)
- x=self.fc3(x)
- return x
-
- if __name__ == '__main__':
- model = CNNNet().to(device)
-
- criterion = nn.CrossEntropyLoss()
- optimizer = optim.Adam(params=model.parameters(), lr=lr)
- # optimizer=optim.SGD(params=model.parameters(),lr=0.001,momentum=0.9)
-
- running_loss = 0.0
-
- for epoch in range(100):
- print(f"=========epoch{epoch+1}=========")
- total = 0
- correct = 0
- for i, data in enumerate(trainLoader, 0):
- inputs, labels = data
- inputs, labels = inputs.to(device), labels.to(device)
-
- optimizer.zero_grad()
- outputs = model(inputs)
- loss = criterion(outputs, labels)
- loss.backward()
- optimizer.step()
- running_loss += loss.item()
- _, predicted = torch.max(outputs.data, 1)
- total += labels.size(0)
- correct += (predicted == labels).sum().item()
- avg_loss = round(running_loss / len(trainLoader), 3)
- running_loss = 0.0
- acc = 100 * correct / total
- print(f"第{epoch + 1}轮准确率:{acc}%,损失值:{avg_loss}")
-
- correct = 0
- total = 0
- with torch.no_grad():
- for data in testLoader:
- inputs, labels = data
- inputs, labels = inputs.to(device), labels.to(device)
-
- outputs = model(inputs)
- _, predicted = torch.max(outputs.data, 1)
- total += labels.size(0)
- correct += (predicted == labels).sum().item()
- print(f"测试集准确率:{100 * correct / total}")
|