import torch
from torch import optim, nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn.functional as Fdevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')#数据加载
transform_train = transforms.Compose([transforms.RandomCrop(32, padding=4),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])transform_test = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4)print(f"Train samples: {len(train_dataset)}")
print(f"Test samples: {len(test_dataset)}")
print(f"Classes: {train_dataset.classes}")#定义MYVGG模型
class MYVGG(nn.Module):def __init__(self, num_classes=10):super(MYVGG, self).__init__()self.features = nn.Sequential(# Block 1nn.Conv2d(3, 64, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.Conv2d(64, 64, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.MaxPool2d(2, 2),# Block 2nn.Conv2d(64, 128, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.Conv2d(128, 128, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.MaxPool2d(2, 2),# Block 3nn.Conv2d(128, 256, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.Conv2d(256, 256, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.MaxPool2d(2, 2),# Block 4nn.Conv2d(256, 512, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.Conv2d(512, 512, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.MaxPool2d(2, 2),# Block 5nn.Conv2d(512, 512, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.Conv2d(512, 512, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.MaxPool2d(2, 2),)self.classifier = nn.Sequential(nn.Dropout(0.5),nn.Linear(512, num_classes))def forward(self, x):x = self.features(x)x = x.view(x.size(0), -1)x = self.classifier(x)return x#训练函数
def train(model, train_loader, epoch_num=50):optimizer = optim.Adam(model.parameters(), lr=0.001)criterion = nn.CrossEntropyLoss()for epoch in range(epoch_num):model.train()for batch_idx, (data, target) in enumerate(train_loader):data, target = data.to(device), target.to(device)optimizer.zero_grad()output = model(data)loss = criterion(output, target)loss.backward()optimizer.step()if batch_idx % 100 == 0:print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(epoch, batch_idx * len(data), len(train_loader.dataset),100. * batch_idx / len(train_loader), loss.item()))torch.save(model.state_dict(), 'cifar10_vgg.pth')print("Model saved to cifar10_vgg.pth")#测试函数
def test(model, test_loader):model.load_state_dict(torch.load('cifar10_vgg.pth', map_location=device))model.eval()correct = 0total = 0with torch.no_grad():for data, target in test_loader:data, target = data.to(device), target.to(device)outputs = model(data)predicted = torch.argmax(outputs, dim=1)total += target.size(0)correct += (predicted == target).sum().item()acc = 100.0 * correct / totalprint(f'Accuracy on CIFAR-10 test set: {acc:.2f}% ({correct}/{total})')#主函数入口
if __name__ == '__main__':model = MYVGG(num_classes=10).to(device)if device == 'cuda':print(f"Using {torch.cuda.device_count()} GPUs")else:print("Using CPU")print("Start training on CIFAR-10...")train(model, train_loader,epoch_num=20)print("Start testing...")test(model, test_loader)print("24信计2班 曾向嵩 2024310143117")
