import math

import torch.nn as nn
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor, Normalize, Pad

from horch.datasets import train_test_split
from horch.models.layers import Conv2d
from horch.nn import Flatten
from horch.train import manual_seed
from horch.optim.lr_scheduler import CosineAnnealingLR

from horch.train.cls import Trainer
from horch.train.metrics import TrainLoss, Loss
from horch.train.cls.metrics import Accuracy

manual_seed(0)


class LeNet5(nn.Sequential):

    def __init__(self):
        super().__init__(
            Conv2d(1, 6, kernel_size=5, norm='default', act='default'),
            nn.AvgPool2d(kernel_size=2, stride=2),
            Conv2d(6, 16, kernel_size=5, norm='default', act='default'),
            nn.AvgPool2d(kernel_size=2, stride=2),
            Flatten(),
            nn.Linear(8 * 8 * 16, 120),
            nn.Linear(120, 84),
            nn.Linear(84, 10),
        )


train_transform = Compose([
    Pad(2),
    ToTensor(),
    Normalize((0.1307,), (0.3081,)),
])

test_transform = Compose([
    Pad(2),
    ToTensor(),
    Normalize((0.1307,), (0.3081,)),
])

data_home = "/Users/hrvvi/Code/study/pytorch/datasets"
ds = MNIST(data_home, train=True, download=False)
ds = train_test_split(ds, test_ratio=0.1, shuffle=True)[1]
ds_train, ds_val = train_test_split(
    ds, test_ratio=0.05, shuffle=True,
    transform=train_transform,
    test_transform=test_transform,
)
ds_test = MNIST(data_home, train=False, download=False, transform=test_transform)
ds_test = train_test_split(ds_test, test_ratio=0.1, shuffle=True)[1]

mul = 1
batch_size = 128
steps_per_epoch = math.ceil(len(ds_train) / batch_size)
min_lr = 0.01 * mul
max_lr = 0.1 * mul

net = LeNet5()
criterion = nn.CrossEntropyLoss()
optimizer = SGD(net.parameters(), lr=min_lr, momentum=0.9, weight_decay=1e-4, nesterov=True)
# lr_scheduler = MultiStepLR(optimizer, [10, 20], gamma=0.1)
lr_scheduler = CosineAnnealingLR(optimizer, T_max=30, eta_min=0.001, warmup=5, warmup_eta_min=0.001)

metrics = {
    'loss': TrainLoss(),
    'acc': Accuracy(),
}

test_metrics = {
    'loss': Loss(criterion),
    'acc': Accuracy(),
}

trainer = Trainer(net, criterion, optimizer, lr_scheduler,
                  metrics=metrics, test_metrics=test_metrics,
                  work_dir="./checkpoints/MNIST-LeNet5")
# summary(net, (1, 32, 32))

train_loader = DataLoader(ds_train, batch_size=128, shuffle=True, num_workers=2, pin_memory=True)
test_loader = DataLoader(ds_test, batch_size=128)
val_loader = DataLoader(ds_val, batch_size=128)

print(trainer.fit(train_loader, 5, val_loader=val_loader))
# trainer.save()
# trainer.load()
# print(trainer.fit(train_loader, 2, val_loader=val_loader))

print(trainer.evaluate(test_loader))
df9e4733bd9ee1f6350fe1139cba6b99ed9722aa
<mask token>


class LeNet5(nn.Sequential):
    <mask token>


<mask token>

<mask token>
manual_seed(0)


class LeNet5(nn.Sequential):

    def __init__(self):
        super().__init__(Conv2d(1, 6, kernel_size=5, norm='default', act=
            'default'), nn.AvgPool2d(kernel_size=2, stride=2), Conv2d(6, 16,
            kernel_size=5, norm='default', act='default'), nn.AvgPool2d(
            kernel_size=2, stride=2), Flatten(), nn.Linear(8 * 8 * 16, 120),
            nn.Linear(120, 84), nn.Linear(84, 10))


<mask token>
print(trainer.fit(train_loader, 5, val_loader=val_loader))
print(trainer.evaluate(test_loader))

<mask token>
manual_seed(0)


class LeNet5(nn.Sequential):

    def __init__(self):
        super().__init__(Conv2d(1, 6, kernel_size=5, norm='default', act=
            'default'), nn.AvgPool2d(kernel_size=2, stride=2), Conv2d(6, 16,
            kernel_size=5, norm='default', act='default'), nn.AvgPool2d(
            kernel_size=2, stride=2), Flatten(), nn.Linear(8 * 8 * 16, 120),
            nn.Linear(120, 84), nn.Linear(84, 10))


train_transform = Compose([Pad(2), ToTensor(), Normalize((0.1307,), (0.3081,))]
    )
test_transform = Compose([Pad(2), ToTensor(), Normalize((0.1307,), (0.3081,))])
data_home = '/Users/hrvvi/Code/study/pytorch/datasets'
ds = MNIST(data_home, train=True, download=False)
ds = train_test_split(ds, test_ratio=0.1, shuffle=True)[1]
ds_train, ds_val = train_test_split(ds, test_ratio=0.05, shuffle=True,
    transform=train_transform, test_transform=test_transform)
ds_test = MNIST(data_home, train=False, download=False, transform=
    test_transform)
ds_test = train_test_split(ds_test, test_ratio=0.1, shuffle=True)[1]
mul = 1
batch_size = 128
steps_per_epoch = math.ceil(len(ds_train) / batch_size)
min_lr = 0.01 * mul
max_lr = 0.1 * mul
net = LeNet5()
criterion = nn.CrossEntropyLoss()
optimizer = SGD(net.parameters(), lr=min_lr, momentum=0.9, weight_decay=
    0.0001, nesterov=True)
lr_scheduler = CosineAnnealingLR(optimizer, T_max=30, eta_min=0.001, warmup
    =5, warmup_eta_min=0.001)
metrics = {'loss': TrainLoss(), 'acc': Accuracy()}
test_metrics = {'loss': Loss(criterion), 'acc': Accuracy()}
trainer = Trainer(net, criterion, optimizer, lr_scheduler, metrics=metrics,
    test_metrics=test_metrics, work_dir='./checkpoints/MNIST-LeNet5')
train_loader = DataLoader(ds_train, batch_size=128, shuffle=True,
    num_workers=2, pin_memory=True)
test_loader = DataLoader(ds_test, batch_size=128)
val_loader = DataLoader(ds_val, batch_size=128)
print(trainer.fit(train_loader, 5, val_loader=val_loader))
print(trainer.evaluate(test_loader))

import math
import torch.nn as nn
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor, Normalize, Pad
from horch.datasets import train_test_split
from horch.models.layers import Conv2d
from horch.nn import Flatten
from horch.train import manual_seed
from horch.optim.lr_scheduler import CosineAnnealingLR
from horch.train.cls import Trainer
from horch.train.metrics import TrainLoss, Loss
from horch.train.cls.metrics import Accuracy
manual_seed(0)


class LeNet5(nn.Sequential):

    def __init__(self):
        super().__init__(Conv2d(1, 6, kernel_size=5, norm='default', act=
            'default'), nn.AvgPool2d(kernel_size=2, stride=2), Conv2d(6, 16,
            kernel_size=5, norm='default', act='default'), nn.AvgPool2d(
            kernel_size=2, stride=2), Flatten(), nn.Linear(8 * 8 * 16, 120),
            nn.Linear(120, 84), nn.Linear(84, 10))


train_transform = Compose([Pad(2), ToTensor(), Normalize((0.1307,), (0.3081,))]
    )
test_transform = Compose([Pad(2), ToTensor(), Normalize((0.1307,), (0.3081,))])
data_home = '/Users/hrvvi/Code/study/pytorch/datasets'
ds = MNIST(data_home, train=True, download=False)
ds = train_test_split(ds, test_ratio=0.1, shuffle=True)[1]
ds_train, ds_val = train_test_split(ds, test_ratio=0.05, shuffle=True,
    transform=train_transform, test_transform=test_transform)
ds_test = MNIST(data_home, train=False, download=False, transform=
    test_transform)
ds_test = train_test_split(ds_test, test_ratio=0.1, shuffle=True)[1]
mul = 1
batch_size = 128
steps_per_epoch = math.ceil(len(ds_train) / batch_size)
min_lr = 0.01 * mul
max_lr = 0.1 * mul
net = LeNet5()
criterion = nn.CrossEntropyLoss()
optimizer = SGD(net.parameters(), lr=min_lr, momentum=0.9, weight_decay=
    0.0001, nesterov=True)
lr_scheduler = CosineAnnealingLR(optimizer, T_max=30, eta_min=0.001, warmup
    =5, warmup_eta_min=0.001)
metrics = {'loss': TrainLoss(), 'acc': Accuracy()}
test_metrics = {'loss': Loss(criterion), 'acc': Accuracy()}
trainer = Trainer(net, criterion, optimizer, lr_scheduler, metrics=metrics,
    test_metrics=test_metrics, work_dir='./checkpoints/MNIST-LeNet5')
train_loader = DataLoader(ds_train, batch_size=128, shuffle=True,
    num_workers=2, pin_memory=True)
test_loader = DataLoader(ds_test, batch_size=128)
val_loader = DataLoader(ds_val, batch_size=128)
print(trainer.fit(train_loader, 5, val_loader=val_loader))
print(trainer.evaluate(test_loader))

import math

import torch.nn as nn
from torch.optim import SGD
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor, Normalize, Pad

from horch.datasets import train_test_split
from horch.models.layers import Conv2d
from horch.nn import Flatten
from horch.train import manual_seed
from horch.optim.lr_scheduler import CosineAnnealingLR

from horch.train.cls import Trainer
from horch.train.metrics import TrainLoss, Loss
from horch.train.cls.metrics import Accuracy

manual_seed(0)


class LeNet5(nn.Sequential):

    def __init__(self):
        super().__init__(
            Conv2d(1, 6, kernel_size=5, norm='default', act='default'),
            nn.AvgPool2d(kernel_size=2, stride=2),
            Conv2d(6, 16, kernel_size=5, norm='default', act='default'),
            nn.AvgPool2d(kernel_size=2, stride=2),
            Flatten(),
            nn.Linear(8 * 8 * 16, 120),
            nn.Linear(120, 84),
            nn.Linear(84, 10),
        )


train_transform = Compose([
    Pad(2),
    ToTensor(),
    Normalize((0.1307,), (0.3081,)),
])

test_transform = Compose([
    Pad(2),
    ToTensor(),
    Normalize((0.1307,), (0.3081,)),
])

data_home = "/Users/hrvvi/Code/study/pytorch/datasets"
ds = MNIST(data_home, train=True, download=False)
ds = train_test_split(ds, test_ratio=0.1, shuffle=True)[1]
ds_train, ds_val = train_test_split(
    ds, test_ratio=0.05, shuffle=True,
    transform=train_transform,
    test_transform=test_transform,
)
ds_test = MNIST(data_home, train=False, download=False, transform=test_transform)
ds_test = train_test_split(ds_test, test_ratio=0.1, shuffle=True)[1]

mul = 1
batch_size = 128
steps_per_epoch = math.ceil(len(ds_train) / batch_size)
min_lr = 0.01 * mul
max_lr = 0.1 * mul

net = LeNet5()
criterion = nn.CrossEntropyLoss()
optimizer = SGD(net.parameters(), lr=min_lr, momentum=0.9, weight_decay=1e-4, nesterov=True)
# lr_scheduler = MultiStepLR(optimizer, [10, 20], gamma=0.1)
lr_scheduler = CosineAnnealingLR(optimizer, T_max=30, eta_min=0.001, warmup=5, warmup_eta_min=0.001)

metrics = {
    'loss': TrainLoss(),
    'acc': Accuracy(),
}

test_metrics = {
    'loss': Loss(criterion),
    'acc': Accuracy(),
}

trainer = Trainer(net, criterion, optimizer, lr_scheduler,
                  metrics=metrics, test_metrics=test_metrics,
                  work_dir="./checkpoints/MNIST-LeNet5")
# summary(net, (1, 32, 32))

train_loader = DataLoader(ds_train, batch_size=128, shuffle=True, num_workers=2, pin_memory=True)
test_loader = DataLoader(ds_test, batch_size=128)
val_loader = DataLoader(ds_val, batch_size=128)

print(trainer.fit(train_loader, 5, val_loader=val_loader))
# trainer.save()
# trainer.load()
# print(trainer.fit(train_loader, 2, val_loader=val_loader))

print(trainer.evaluate(test_loader))