|
- 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, padding=4),
- 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=64, shuffle=True, num_workers=4, pin_memory=True)
- testLoader = DataLoader(testSet, batch_size=64, shuffle=False, num_workers=4, pin_memory=True)
-
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-
- class VGG(nn.Module):
- def __init__(self):
- super(VGG, self).__init__()
- self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
- self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
- self.pool1 = nn.MaxPool2d((2, 2), padding=1)
- self.bn1 = nn.BatchNorm2d(64)
- self.relu1 = nn.ReLU()
-
- self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
- self.conv4 = nn.Conv2d(128, 128, 3, padding=1)
- self.pool2 = nn.MaxPool2d((2, 2), padding=1)
- self.bn2 = nn.BatchNorm2d(128)
- self.relu2 = nn.ReLU()
-
- self.conv5 = nn.Conv2d(128, 256, 3, padding=1)
- self.conv6 = nn.Conv2d(256, 256, 3, padding=1)
- self.conv7 = nn.Conv2d(256, 256, 3, padding=1)
- self.pool3 = nn.MaxPool2d((2, 2), padding=1)
- self.bn3 = nn.BatchNorm2d(256)
- self.relu3 = nn.ReLU()
-
- self.conv8 = nn.Conv2d(256, 512, 3, padding=1)
- self.conv9 = nn.Conv2d(512, 512, 3, padding=1)
- self.conv10 = nn.Conv2d(512, 512, 3, padding=1)
- self.pool4 = nn.MaxPool2d((2, 2), padding=1)
- self.bn4 = nn.BatchNorm2d(512)
- self.relu4 = nn.ReLU()
-
- self.conv11 = nn.Conv2d(512, 512, 3, padding=1)
- self.conv12 = nn.Conv2d(512, 512, 3, padding=1)
- self.conv13 = nn.Conv2d(512, 512, 3, padding=1)
- self.pool5 = nn.MaxPool2d((2, 2), padding=1)
- self.bn5 = nn.BatchNorm2d(512)
- self.relu5 = nn.ReLU()
-
- self.fc14 = nn.Linear(512 * 2 * 2, 4096)
- self.relu6 = nn.ReLU()
- self.drop1 = nn.Dropout2d()
- self.fc15 = nn.Linear(4096, 4096)
- self.relu7 = nn.ReLU()
- self.drop2 = nn.Dropout2d()
- self.fc16 = nn.Linear(4096, 10)
-
- def forward(self, x):
- # train = transforms.Resize((224, 224))
- # x = train(x)
- x = self.conv1(x)
- x = self.conv2(x)
- x = self.pool1(x)
- x = self.bn1(x)
- x = self.relu1(x)
-
- x = self.conv3(x)
- x = self.conv4(x)
- x = self.pool2(x)
- x = self.bn2(x)
- x = self.relu2(x)
-
- x = self.conv5(x)
- x = self.conv6(x)
- x = self.conv7(x)
- x = self.pool3(x)
- x = self.bn3(x)
- x = self.relu3(x)
-
- x = self.conv8(x)
- x = self.conv9(x)
- x = self.conv10(x)
- x = self.pool4(x)
- x = self.bn4(x)
- x = self.relu4(x)
-
- x = self.conv11(x)
- x = self.conv12(x)
- x = self.conv13(x)
- x = self.pool5(x)
- x = self.bn5(x)
- x = self.relu5(x)
- # print(x.shape)
-
- x = x.view(-1, 512 * 2 * 2)
-
- x = self.fc14(x)
- x = self.relu6(x)
- x = self.drop1(x)
- x = self.fc15(x)
- x = self.relu7(x)
- x = self.drop2(x)
- x = self.fc16(x)
- return x
-
-
- if __name__ == '__main__':
- model = VGG().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(50):
- total = 0
- correct = 0
- print(f"=========epoch{epoch + 1}=========")
-
- 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()
- # print(i)
- # if i % 2000 == 1999:
- # # print(loss)
- # print(f"[{epoch+1},{i+1}] loss: {round(running_loss/2000,3)}")
- # running_loss = 0.0
- avg_loss = round(running_loss / len(trainLoader), 2)
- running_loss = 0.0
- acc = 100 * correct / total
- print(f"第{epoch + 1}轮准确率:{acc}%,损失值:{avg_loss}")
- torch.save(model.state_dict(),"/model/model_vgg_cifar10.pt")
-
- total = 0
- correct = 0
- running_loss = 0.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}%")
|