import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm

from torchvision import transforms
import torchvision

from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.ToTensor(),
])

batch_size = 4
train_set = torchvision.datasets.ImageFolder(root='data/cat_vs_dog/train', transform=transform)


train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True,
                          num_workers=0)  # Batch Size定义：一次训练所选取的样本数。 Batch Size的大小影响模型的优化程度和速度。

test_dataset = torchvision.datasets.ImageFolder(root='data/cat_vs_dog/test', transform=transform)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True,
                         num_workers=0)  # Batch Size定义：一次训练所选取的样本数。 Batch Size的大小影响模型的优化程度和速度。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

net = torchvision.models.resnet50(weights=True)
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs, 2)  # 将输出维度修改为2

criterion = nn.CrossEntropyLoss()
net = net.to(device)
optimizer = torch.optim.Adam(lr=0.0001, params=net.parameters())
eposhs = 100

for epoch in range(eposhs):
    correct = 0
    sum_loss = 0
    total = 0
    for inputs, labels in tqdm(train_loader):
        inputs = inputs.to(device)
        labels = labels.to(device)

        output = net(inputs)
        loss = criterion(output, labels)
        total = total + labels.size(0)
        optimizer.zero_grad()
        _, predicted = torch.max(output.data, 1)
        loss.backward()
        optimizer.step()
        correct = correct + (predicted == labels).sum().item()

    acc = correct / total
    print('准确率是:', acc)