import torch
import torch.nn as nn
from torchvision import datasets, transforms
import autoencoder


def train():
    transform = transforms.ToTensor()
    mnist_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    data_loader = torch.utils.data.DataLoader(dataset=mnist_data,
                                              batch_size=64,
                                              shuffle=True)

    model = autoencoder.Autoencoder_CNN()
    model.train()
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)

    num_epochs = 3
    outputs = []
    for epoch in range(num_epochs):
        for (img, _) in data_loader:
            recon = model(img)
            loss = criterion(recon, img)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(f'Epoch:{epoch + 1}, Loss:{loss.item():.4f}')
        outputs.append((epoch, img, recon))

    torch.save(model.state_dict(), 'models/cnn.h5')