Alexnet模型以及训练代码

import numpy as np import torch import torchvision import argparse from torch import nn from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from torchvision import transforms from torchvision.transforms import InterpolationMode from alexnet import AlexNet import ssl # 全局取消证书验证 ssl._create_default_https_context = ssl._create_unverified_context device= torch.device('cuda' if torch.cuda.is_available() else 'cpu')# print(device) def train(args): # 超参数设置，方便管理 num_epochs = args.max_epoch batch_size = args.batch_size learning_rate = args.lr image_size = args.image_size momentum = args.momentum # 设置数据集的格式 transform = transforms.Compose([transforms.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC), transforms.ToTensor(), transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.247, 0.2435, 0.2616]) ]) # 数据加载 # 如果没有这个数据集的话会自动下载 train_data = torchvision.datasets.CIFAR10(root="dataset",download=True,transform=transform,train=True) test_data = torchvision.datasets.CIFAR10(root="dataset",download=True,transform=transform,train=False) train_dataloader = DataLoader(train_data, batch_size=batch_size) test_dataloader = DataLoader(test_data, batch_size=batch_size) print('Dataload is Ready') # 添加tensorboard路径 writer = SummaryWriter(log_dir=args.SummerWriter_log) # 模型加载 model = AlexNet(num_classes=10).to(device) # 参数量估计 total = sum([param.nelement() for param in model.parameters()]) print("Number of parameters: %.2fM" % (total / 1e6)) # Loss and optimizer # 选择交叉熵作为损失函数 criterion = nn.CrossEntropyLoss() # 选择SGD为优化器 optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum) total_train_step = 0#记录训练次数 total_test_step=0#记录测试次数 # 开始训练 for epoch in range(num_epochs): print("---------------第{}轮训练开始-------------".format(epoch + 1)) for i, (images, labels) in enumerate(train_dataloader): images = images.to(device) labels = labels.to(device) # Forward pass outputs = model(images) loss = criterion(outputs, labels) # Backward and optimize optimizer.zero_grad() loss.backward() optimizer.step() total_train_step = total_train_step + 1 if (i + 1) % args.print_map_epoch == 0:# 100次显示一次loss print("Epoch [{}/{}], Step [{}] Loss: {:.4f}" .format(epoch + 1, num_epochs, total_train_step, loss.item())) writer.add_scalar("train_loss", loss.item(), total_train_step) # Test the model model.eval() total_test_loss = 0 total_accuary = 0 # 正确率 with torch.no_grad(): correct = 0 total = 0 for images, labels in test_dataloader: images = images.to(device) labels = labels.to(device) outputs = model(images) loss = criterion(outputs, labels) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() total_test_loss += loss total_accuary += correct print('Accuracy of the model on the test images: {} %'.format(100 * correct / total)) writer.add_scalar("test_loss",total_test_loss,total_test_step) writer.add_scalar("test_accuary", correct / total, total_test_step) total_test_step += 1 # Save the model checkpoint torch.save(model, 'weights.pth') if __name__ == "__main__": parser = argparse.ArgumentParser() '''------------------------------------------调 节 部 分------------------------------------------------------''' parser.add_argument('--max_epoch', type=int, default=40, help='total epoch') parser.add_argument('--device_num', type=str, default='cpu', help='select GPU or cpu') parser.add_argument('--image_size', type=int, default=224, help='if crop img, img will be resized to the size') parser.add_argument('--batch_size', type=int, default=32, help='batch size, recommended 16') parser.add_argument('--lr', type=float, default=0.001, help='learning rate') parser.add_argument('--momentum', type=float, default=0.90, help='choice a float in 0-1') parser.add_argument('--print_map_epoch', type=int, default=100, help='') parser.add_argument('--SummerWriter_log', type=str, default='Alexnet', help='') args = parser.parse_args() train(args)