# file     : train.py

# time     : 2022/8/11 上午10:03

# function :

import torchvision.datasets

from model import *

from torch.utils.data import DataLoader

# DataSet

train_data = torchvision.datasets.CIFAR10("../dataset", train=True, transform=torchvision.transforms.ToTensor(), download=False)

test_data = torchvision.datasets.CIFAR10("../dataset", train=False, transform=torchvision.transforms.ToTensor(), download=False)

# len长度

train_data_size = len(train_data)

test_data_size = len(test_data)

print(format(train_data_size))

# ctrl+D

print(format(test_data_size))

# liyong DataLoader

train_dataloader = DataLoader(train_data, batch_size=64)

test_dataloader = DataLoader(test_data, batch_size=64)

# create

tudui = Tudui()

#loss

loss_fn = nn.CrossEntropyLoss()

#

# learning_rate = 0.01

# 1e-2 = 1 x (10)^(-2) = 1/100 = 0.01

learning_rate = 1e-2

optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate)

# shezhi

total_train_step = 0

#

total_test_step = 0

#

epoch = 10

for i in range(epoch):

    print("------------第{} 轮训练开始----------- ".format(i+1))

    #

    for data in train_dataloader:

        imgs, targets = data

        output = tudui(imgs)

        loss = loss_fn(output, targets)

        #

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()

        total_train_step = total_train_step + 1

        print("训练次数：{}, loss:{}".format(total_train_step, loss.item))