__all__ = ["Compose", "ToTensor", "ToCVImage", "Normalize", "Resize", "Scale", "CenterCrop", "Pad",
           "Lambda", "RandomApply", "RandomChoice", "RandomOrder", "RandomCrop", "RandomHorizontalFlip",
           "RandomVerticalFlip", "RandomResizedCrop", "RandomSizedCrop", "FiveCrop", "TenCrop", "LinearTransformation",
           "ColorJitter", "RandomRotation", "RandomAffine", "Grayscale", "RandomGrayscale",
           "RandomPerspective", "RandomErasing", "RandomGaussianNoise", "RandomPoissonNoise", "RandomSPNoise",
           "RandomTransforms", "Cutout"]

import torch
from torchtoolbox.tools import summary
from torchvision.models.mobilenet import mobilenet_v2
model = mobilenet_v2()
summary(model, torch.rand((1, 3, 224, 224)))

from torchtoolbox import metric

# define first
top1_acc = metric.Accuracy(name='Top1 Accuracy')
top5_acc = metric.TopKAccuracy(top=5, name='Top5 Accuracy')
loss_record = metric.NumericalCost(name='Loss')

# reset before using
top1_acc.reset()
top5_acc.reset()
loss_record.reset()

...
model.eval()
for data, labels in val_data:
    data = data.to(device, non_blocking=True)
    labels = labels.to(device, non_blocking=True)

    outputs = model(data)
    losses = Loss(outputs, labels)
    # update/record
    top1_acc.update(outputs, labels)
    top5_acc.update(outputs, labels)
    loss_record.update(losses)
    test_msg = 'Test Epoch {}: {}:{:.5}, {}:{:.5}, {}:{:.5}\n'.format(
    epoch, top1_acc.name, top1_acc.get(), top5_acc.name, top5_acc.get(),
    loss_record.name, loss_record.get())

print(test_msg)

from torchtoolbox.nn.init import KaimingInitializer

model = XXX
initializer = KaimingInitializer()
model.apply(initializer)

from torch import nn
from torchtoolbox.nn import AdaptiveSequential
import torch


class n_to_n(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 3, 1, 1, bias=False)
        self.conv2 = nn.Conv2d(3, 3, 1, 1, bias=False)

    def forward(self, x1, x2):
        y1 = self.conv1(x1)
        y2 = self.conv2(x2)
        return y1, y2


class n_to_one(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 3, 1, 1, bias=False)
        self.conv2 = nn.Conv2d(3, 3, 1, 1, bias=False)

    def forward(self, x1, x2):
        y1 = self.conv1(x1)
        y2 = self.conv2(x2)
        return y1 + y2


class one_to_n(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 3, 1, 1, bias=False)
        self.conv2 = nn.Conv2d(3, 3, 1, 1, bias=False)

    def forward(self, x):
        y1 = self.conv1(x)
        y2 = self.conv2(x)
        return y1, y2

seq = AdaptiveSequential(one_to_n(), n_to_n(), n_to_one()).cuda()
td = torch.rand(1, 3, 32, 32).cuda()

out = seq(td)
print(out.size())

# output
# torch.Size([1, 3, 32, 32])

from torchtoolbox.tools.convert_lmdb import generate_lmdb_dataset, raw_reader
from torchvision.datasets import ImageFolder

dt = ImageFolder(..., loader=raw_reader)
save_dir = XXX 
dataset_name = YYY
generate_lmdb_dataset(dt, save_dir=save_dir, name=dataset_name)

from torchtoolbox.data import ImageLMDB

dt = ImageLMDB(db_path=save_dir, db_name=dataset_name, ...)

from torchtoolbox.nn import Activation

relu = Activation('relu', auto_optimize=True)

from torchtoolbox.nn import LabelSmoothingLoss
# The num classes of your task should be defined.
classes = 10
# Loss
Loss = LabelSmoothingLoss(classes, smoothing=0.1)

...
for i, (data, labels) in enumerate(train_data):
    data = data.to(device, non_blocking=True)
    labels = labels.to(device, non_blocking=True)

    optimizer.zero_grad()
    outputs = model(data)
    # just use as usual.
    loss = Loss(outputs, labels)
    loss.backward()
    optimizer.step()

from torchtoolbox.optimizer import CosineWarmupLr

optimizer = optim.SGD(...)
# define scheduler
# `batches_pre_epoch` means how many batches(times update/step the model) within one epoch.
# `warmup_epochs` means increase lr how many epochs to `base_lr`.
# you can find more details in file.
lr_scheduler = CosineWarmupLr(optimizer, batches_pre_epoch, epochs,
                              base_lr=lr, warmup_epochs=warmup_epochs)
...
for i, (data, labels) in enumerate(train_data):
    ...
    optimizer.step()
    # remember to step/update status here.
    lr_scheduler.step()
    ...

from torchtoolbox.nn import SwitchNorm2d, SwitchNorm3d

from torchtoolbox.nn import Swish

from torchtoolbox.optimizer import Lookahead
from torch import optim

optimizer = optim.Adam(...)
optimizer = Lookahead(optimizer)

from torchtoolbox.tools import mixup_data, mixup_criterion

# set beta distributed parm, 0.2 is recommend.
alpha = 0.2
for i, (data, labels) in enumerate(train_data):
    data = data.to(device, non_blocking=True)
    labels = labels.to(device, non_blocking=True)

    data, labels_a, labels_b, lam = mixup_data(data, labels, alpha)
    optimizer.zero_grad()
    outputs = model(data)
    loss = mixup_criterion(Loss, outputs, labels_a, labels_b, lam)

    loss.backward()
    optimizer.update()

from torchvision import transforms
from torchtoolbox.transform import Cutout

_train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    Cutout(),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(0.4, 0.4, 0.4),
    transforms.ToTensor(),
    normalize,
])

from torchtoolbox.tools import no_decay_bias
from torch import optim

model = XXX
parameters = no_decay_bias(model)
optimizer = optim.SGD(parameters, ...)

from torchtoolbox.nn.loss import ArcLoss, CosLoss, L2Softmax

from torchtoolbox.nn import Swish, HardSwish, HardSigmoid

swish = Swish()
hswish = HardSwish()
hsigmoid = HardSigmoid()

from torchtoolbox.nn.init import ZeroLastGamma

model == XXX
init = ZeroLastGamma(block_name='Bottleneck', bn_name='bn3')
model.apply(init)

from torchtoolbox.optimizer import SGD_GC

optimizer = SGD_GC(model.parameters())
optimizer.zero_grad()
loss.backward()
optimizer.step()

# for ImageNet
from torchtoolbox.transform import ImageNetPolicy, Compose, \
            RandomResizedCrop, RandomHorizontalFlip, ToTensor

transforms = Compose([RandomResizedCrop(224), 
                      RandomHorizontalFlip(),
                      ImageNetPolicy,
                      ToTensor()])