import os.path

import torch
import torch.nn as nn
from PIL import Image
from torchvision import datasets

from HerbModel import Model

from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import Dataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.transforms import transforms

img_transforms =transforms.Compose([


   transforms.Resize(128),
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(128),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5],[0.5, 0.5, 0.5])
])

img_transforms2 =transforms.Compose([

    transforms.Resize(32),
    # transforms.RandomHorizontalFlip(),
    # transforms.CenterCrop(32),
    transforms.ToTensor(),
    transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])
])


# 数据加载
data_dir = '.'
train_data = datasets.ImageFolder(os.path.join(data_dir, 'yaocai'), img_transforms)

train_loader = DataLoader(train_data, 4,shuffle=True)

print("训练数据集数量：{}".format(len(train_data)))
model = Model()
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
writer = SummaryWriter("yc_logs")
epoch = 50
step = 0
for i in range(epoch):
    print("第{}轮训练开始。。。。。".format(i+1))
    for data in train_loader:
        imgs, target = data
        # print(imgs.shape)
        output = model(imgs)
        # print(output)


        loss = loss_fn(output, target)

        print("损失：{}".format(loss))

        argmax = output.argmax(1)
        acc = (argmax == target).sum()


        print("结果：{}".format(argmax))
        print("实际：{}".format(target))
        print("正确率：{}".format(acc/len(imgs)*100))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        step = step + 1
        print("第{}次训练".format(step))

step = 0

writer.close()