AlexNet卷积神经网络图像分类Pytorch训练代码 使用Cifar100数据集

# 训练代码，训练集使用CIFAR100，初次训练会自动下载 import torch import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms from torch.optim.lr_scheduler import StepLR from net import AlexNet # 训练参数 batch_size = 8 epochs = 10 lr = 0.01 gamma = 0.7 no_cuda = False seed = 1 log_interval = 10 save_model = True def train(model, device, train_loader, optimizer, epoch, log_interval): 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 = F.cross_entropy(output, target) loss.backward() optimizer.step() if batch_idx % log_interval == 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())) if save_model: torch.save(model.state_dict(), "model.pth") def main(): use_cuda = not no_cuda and torch.cuda.is_available() torch.manual_seed(seed) device = torch.device("cuda" if use_cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} train_loader = torch.utils.data.DataLoader( datasets.CIFAR100('./data', train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Resize(224), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ])), batch_size=batch_size, shuffle=True, **kwargs) model = AlexNet(num_labels=100).to(device) optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9) scheduler = StepLR(optimizer, step_size=1, gamma=gamma) for epoch in range(1, epochs + 1): train(model, device, train_loader, optimizer, epoch, log_interval) scheduler.step() if __name__ == '__main__': main()