
from model import *
import torchvision
import torch
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from torch import nn
from torch.utils.data import DataLoader
data_transforms = transforms.Compose([
transforms.RandomRotation(45),
transforms.ToTensor(),
])
train_data = torchvision.datasets.CIFAR10(root="D:\pythonProject_pytorchstudy", train=True, transform=data_transforms, download=False)
test_data = torchvision.datasets.CIFAR10(root="D:\pythonProject_pytorchstudy", train=False, transform=torchvision.transforms.ToTensor(), download=False)
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练集的长度为:{}".format(train_data_size))
print("测试集的长度为:{}".format(test_data_size))
train_dataloader =DataLoader(train_data,batch_size=64)
test_dataloader =DataLoader(test_data,batch_size=64)
Yolo = My_Model()
if torch.cuda.is_available():
Yolo = My_Model().cuda()
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
loss_fn = loss_fn.cuda()
learning_rate = 0.01
optimizer = torch.optim.SGD(Yolo.parameters(), lr = learning_rate, )
total_train_step = 0
total_test_step = 0
epoch = 10
writer = SummaryWriter("D:\pythonProject_pytorchstudy\cifar-10-batches-py\logs_train")
for i in range(epoch):
print("第{}轮训练开始".format(i+1))
Yolo.train()
for data in train_dataloader:
imgs,targets = data
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = Yolo(imgs)
loss = loss_fn(outputs,targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_step += 1
if total_train_step % 30 ==0:
print("Iteration:{},loss:{}".format(total_train_step,loss.item()))
writer.add_scalar("train_loss", loss.item(),total_train_step)
Yolo.eval()
total_test_loss = 0
total_accuracy = 0
with torch.no_grad():
for data in test_dataloader:
imgs, targets = data
if torch.cuda.is_available():
imgs = imgs.cuda()
targets = targets.cuda()
outputs = Yolo(imgs)
loss = loss_fn(outputs,targets)
total_test_loss = total_test_loss + loss.item()
accuracy = (outputs.argmax(1) == targets).sum()
total_accuracy = total_accuracy + accuracy
print("整体测试集上的Loss{}".format(total_test_loss))
print("整体测试集上的正确率:{}".format(total_accuracy/test_data_size))
writer.add_scalar("test_loss",total_test_loss,total_test_step)
writer.add_scalar("test_accuracy",total_accuracy/test_data_size,total_test_step)
total_train_step += 1
torch.save(Yolo,"YOLO_{}".format(i+1))
print("模型已保存")
writer.close()
import torch
from torch import nn
class My_Model(nn.Module):
def __init__(self):
super(My_Model, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64 * 4 * 4, 64),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x