赞
踩
- train_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=True, transform=torchvision.transforms.ToTensor())
- test_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=False, transform=torchvision.transforms.ToTensor())
- import torch
- import torchvision
- from torch import nn
- from torch.utils.data import DataLoader
- from torch.utils.tensorboard import SummaryWriter
- import time
-
-
- train_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=True, transform=torchvision.transforms.ToTensor())
- test_data = torchvision.datasets.CIFAR10("../input/cifar10-python", train=False, transform=torchvision.transforms.ToTensor())
-
- train_dataloader = DataLoader(train_data, batch_size=64, drop_last=True)
- test_dataloader = DataLoader(test_data, batch_size=64, drop_last=True)
- # print(len(train_dataloader)) #781
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-
- test_data_size = len(test_dataloader) * 64
- print(f'测试集大小为:{test_data_size}')
- writer = SummaryWriter("../model_logs")
-
- loss_fn = nn.CrossEntropyLoss(reduction='mean')
- loss_fn = loss_fn.to(device)
- time_able = False # True
-
- class Model(nn.Module):
- def __init__(self):
- super(Model, self).__init__()
- self.model1 = nn.Sequential(
- nn.Conv2d(3, 32, 5, padding=2),
- nn.MaxPool2d(2),
- nn.Conv2d(32, 32, 5, padding=2),
- nn.MaxPool2d(2),
- nn.Conv2d(32, 64, 5, padding=2),
- nn.MaxPool2d(2),
- nn.Flatten(),
- nn.Linear(1024, 64),# 182528
- nn.Linear(64, 10)
- )
-
- def forward(self, x):
- x = self.model1(x)
-
- return x
-
- model = Model()
- model = model.to(device)
- optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
- epoch = 50
- running_loss = 0
- total_train_step = 0
- total_test_step = 0
- if time_able:
- str_time = time.time()
- for i in range(epoch):
- print(f'第{i + 1}次epoch')
- for data in train_dataloader:
- imgs, targets = data
- imgs = imgs.to(device)
- targets = targets.to(device)
- output = model(imgs)
- loss = loss_fn(output, targets)
- optimizer.zero_grad()
- loss.backward()
- optimizer.step()
- total_train_step += 1
- if total_train_step % 100 == 0:
- if time_able:
- end_time = time.time()
- print(f'{str_time-end_time}')
- print(f'第{total_train_step}次训练,loss = {loss.item()}')
- writer.add_scalar("train_loss", loss.item(), total_train_step)
- # 测试
- total_test_loss = 0
- total_accuracy = 0
- with torch.no_grad():
- for data in test_dataloader:
- imgs, targets = data
- imgs = imgs.to(device)
- targets = targets.to(device)
- outputs = model(imgs)
- loss = loss_fn(outputs, targets)
- total_test_loss = total_test_loss + loss
- accuracy = (outputs.argmax(1) == targets).sum()
- total_accuracy += accuracy
- total_test_loss = total_test_loss / test_data_size
- print(f'整体测试集上的loss = {total_test_loss}')
- print(f'整体测试集正确率 = {total_accuracy / test_data_size}')
- writer.add_scalar("test_loss", total_test_loss.item(), total_test_step)
- writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)
- total_test_step += 1
-
- writer.close()
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。