examples/basic_tutorials/cifar10_cnn.py

#! /usr/bin/python
# -*- coding: utf-8 -*-


################################ TensorLayerX and Torch can be mixed programming. #################################

import os
# os.environ['TL_BACKEND'] = 'paddle'
# os.environ['TL_BACKEND'] = 'tensorflow'
# os.environ['TL_BACKEND'] = 'mindspore'
os.environ['TL_BACKEND'] = 'torch'


import time
from tensorlayerx.dataflow import Dataset, DataLoader
from tensorlayerx.vision.transforms import (
    Compose, Resize, RandomFlipHorizontal, RandomContrast, RandomBrightness, StandardizePerImage, RandomCrop
)
from tensorlayerx.model import TrainOneStep
from tensorlayerx.nn import Module
import tensorlayerx as tlx
from tensorlayerx.nn import (Conv2d, Linear, Flatten, MaxPool2d, BatchNorm2d)
# enable debug logging
tlx.logging.set_verbosity(tlx.logging.DEBUG)

# ################## Download and prepare the CIFAR10 dataset ##################
# This is just some way of getting the CIFAR10 dataset from an online location
# and loading it into numpy arrays with shape [32,32,3]
X_train, y_train, X_test, y_test = tlx.files.load_cifar10_dataset(shape=(-1, 32, 32, 3), plotable=False)

# ################## CIFAR10 dataset ##################
# We define a Dataset class for Loading CIFAR10 images and labels.
class make_dataset(Dataset):

    def __init__(self, data, label, transforms):
        self.data = data
        self.label = label
        self.transforms = transforms

    def __getitem__(self, idx):
        x = self.data[idx].astype('uint8')
        y = self.label[idx].astype('int64')
        x = self.transforms(x)

        return x, y

    def __len__(self):

        return len(self.label)

# We define the CIFAR10 iamges preprocessing pipeline.
train_transforms = Compose( # Combining multiple operations sequentially
    [
        RandomCrop(size=[24, 24]), #random crop from images to shape [24, 24]
        RandomFlipHorizontal(), # random invert each image horizontally by probability
        RandomBrightness(brightness_factor=(0.5, 1.5)), # Within the range of values (0.5, 1.5), adjust brightness randomly
        RandomContrast(contrast_factor=(0.5, 1.5)), # Within the range of values (0.5, 1.5), adjust contrast randomly
        StandardizePerImage() #Normalize the values of each image to [-1, 1]
    ]
)

test_transforms = Compose([Resize(size=(24, 24)), StandardizePerImage()])

# We use DataLoader to batch and shuffle data, and make data into iterators.
train_dataset = make_dataset(data=X_train, label=y_train, transforms=train_transforms)
test_dataset = make_dataset(data=X_test, label=y_test, transforms=test_transforms)

train_dataset = DataLoader(train_dataset, batch_size=128, shuffle=True)
test_dataset = DataLoader(test_dataset, batch_size=128)

# ################## CNN network ##################
class CNN(Module):

    def __init__(self):
        super(CNN, self).__init__()
        # Parameter initialization method
        W_init = tlx.nn.initializers.truncated_normal(stddev=5e-2)
        W_init2 = tlx.nn.initializers.truncated_normal(stddev=0.04)
        b_init2 = tlx.nn.initializers.constant(value=0.1)

        # 2D Convolutional Neural Network, Set padding method "SAME", convolutional kernel size [5,5], stride [1,1], in channels, out channels
        self.conv1 = Conv2d(64, (5, 5), (1, 1), padding='SAME', W_init=W_init, b_init=None, name='conv1', in_channels=3)
        # Add 2D BatchNormalize, using ReLU for output.
        self.bn = BatchNorm2d(num_features=64, act=tlx.nn.ReLU)
        # Add 2D Max pooling layer.
        self.maxpool1 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool1')

        self.conv2 = Conv2d(
            64, (5, 5), (1, 1), padding='SAME', act=tlx.nn.ReLU, W_init=W_init, name='conv2', in_channels=64
        )
        self.maxpool2 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool2')
        # Flatten 2D data to 1D data
        self.flatten = Flatten(name='flatten')
        # Linear layer with 384 units, using ReLU for output.
        self.linear1 = Linear(384, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear1relu', in_features=2304)
        self.linear2 = Linear(192, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear2relu', in_features=384)
        self.linear3 = Linear(10, act=None, W_init=W_init2, name='output', in_features=192)

    # We define the forward computation process.
    def forward(self, x):
        z = self.conv1(x)
        z = self.bn(z)
        z = self.maxpool1(z)
        z = self.conv2(z)
        z = self.maxpool2(z)
        z = self.flatten(z)
        z = self.linear1(z)
        z = self.linear2(z)
        z = self.linear3(z)
        return z


# get the network
net = CNN()

# training settings
n_epoch = 500
learning_rate = 0.0001
print_freq = 5
n_step_epoch = int(len(y_train) / 128)
n_step = n_epoch * n_step_epoch
shuffle_buffer_size = 128
# Get training parameters
train_weights = net.trainable_weights
# Define the optimizer, use the Adam optimizer.
optimizer = tlx.optimizers.Adam(learning_rate)
# Define evaluation metrics.
metrics = tlx.metrics.Accuracy()

# Define the loss calculation process
class WithLoss(Module):

    def __init__(self, net, loss_fn):
        super(WithLoss, self).__init__()
        self._net = net
        self._loss_fn = loss_fn

    def forward(self, data, label):
        out = self._net(data)
        loss = self._loss_fn(out, label)
        return loss


net_with_loss = WithLoss(net, loss_fn=tlx.losses.softmax_cross_entropy_with_logits)
# Initialize one-step training
net_with_train = TrainOneStep(net_with_loss, optimizer, train_weights)

# Custom training loops
for epoch in range(n_epoch):
    start_time = time.time()
    # Set the network to training state
    net.set_train()
    train_loss, train_acc, n_iter = 0, 0, 0
    # Get training data and labels
    for X_batch, y_batch in train_dataset:
        # Calculate the loss value, and automatically complete the gradient update
        _loss_ce = net_with_train(X_batch, y_batch)
        train_loss += _loss_ce

        n_iter += 1
        _logits = net(X_batch)
        # Calculate accuracy
        metrics.update(_logits, y_batch)
        train_acc += metrics.result()
        metrics.reset()
        print("Epoch {} of {} took {}".format(epoch + 1, n_epoch, time.time() - start_time))
        print("   train loss: {}".format(train_loss / n_iter))
        print("   train acc:  {}".format(train_acc / n_iter))


################################ TensorLayerX and Jittor can be mixed programming. #################################


# import os
# import time
# import numpy as np
# import tensorlayerx as tlx
# from tensorlayerx.dataflow import Dataset, DataLoader
# from tensorlayerx.vision.transforms import (
#     Compose, Resize, RandomFlipHorizontal, RandomContrast, RandomBrightness, StandardizePerImage, RandomCrop
# )
# from tensorlayerx.nn import Conv2d, Linear, Flatten, Module
# from tensorlayerx.optimizers import Adam
# from tqdm import tqdm

# # Enable debug logging
# tlx.logging.set_verbosity(tlx.logging.DEBUG)

# os.environ['TL_BACKEND'] = 'jittor'


# # Download and prepare the CIFAR10 dataset with progress bar
# print("Downloading CIFAR10 dataset...")
# X_train, y_train, X_test, y_test = tlx.files.load_cifar10_dataset(shape=(-1, 32, 32, 3), plotable=False)

# # Define the CIFAR10 dataset
# class CIFAR10Dataset(Dataset):
#     def __init__(self, data, label, transforms):
#         self.data = data
#         self.label = label
#         self.transforms = transforms

#     def __getitem__(self, idx):
#         x = self.data[idx].astype('uint8')
#         y = self.label[idx].astype('int64')
#         x = self.transforms(x)
#         return x, y

#     def __len__(self):
#         return len(self.label)

# # Define the CIFAR10 images preprocessing pipeline
# train_transforms = Compose([
#     RandomCrop(size=[24, 24]),
#     RandomFlipHorizontal(),
#     RandomBrightness(brightness_factor=(0.5, 1.5)),
#     RandomContrast(contrast_factor=(0.5, 1.5)),
#     StandardizePerImage()
# ])

# test_transforms = Compose([Resize(size=(24, 24)), StandardizePerImage()])

# # Create DataLoaders for training and testing
# print("Processing CIFAR10 dataset...")
# train_dataset = CIFAR10Dataset(data=X_train, label=y_train, transforms=train_transforms)
# test_dataset = CIFAR10Dataset(data=X_test, label=y_test, transforms=test_transforms)

# train_dataloader = DataLoader(train_dataset, batch_size=128, shuffle=True)
# test_dataloader = DataLoader(test_dataset, batch_size=128)

# # Define a simple CNN model
# class SimpleCNN(Module):
#     def __init__(self):
#         super(SimpleCNN, self).__init__()
#         self.conv1 = Conv2d(16, (3, 3), (1, 1), padding='SAME', act=tlx.nn.ReLU, in_channels=3)
#         self.flatten = Flatten()
#         self.fc1 = Linear(out_features=64, act=tlx.nn.ReLU, in_features=16 * 24 * 24)
#         self.fc2 = Linear(out_features=10, act=None, in_features=64)

#     def forward(self, x):
#         z = self.conv1(x)
#         z = self.flatten(z)
#         z = self.fc1(z)
#         z = self.fc2(z)
#         return z

# # Instantiate the model
# model = SimpleCNN()

# # Define the optimizer
# optimizer = Adam(model.trainable_weights, lr=0.001)

# # Define the loss function
# loss_fn = tlx.losses.softmax_cross_entropy_with_logits

# # Training loop
# n_epoch = 2
# for epoch in range(n_epoch):
#     start_time = time.time()
#     model.set_train()
#     train_loss, n_iter = 0, 0

#     with tqdm(total=len(train_dataloader), desc=f"Epoch {epoch + 1}/{n_epoch}", unit="batch") as pbar:
#         for X_batch, y_batch in train_dataloader:
#             X_batch = tlx.convert_to_tensor(X_batch)
#             y_batch = tlx.convert_to_tensor(y_batch)
#             _logits = model(X_batch)
#             loss = loss_fn(_logits, y_batch)
            
#             optimizer.zero_grad()
#             optimizer.step(loss)
            
#             train_loss += loss.item()
#             n_iter += 1
#             pbar.update(1)

#     print(f"Epoch {epoch + 1} of {n_epoch} took {time.time() - start_time:.2f}s")
#     print(f"   train loss: {train_loss / n_iter:.4f}")


################################ TensorLayerX and TensorFlow can be mixed programming. #################################
# import os
# os.environ['TL_BACKEND'] = 'tensorflow'
#
# import time
# import numpy as np
# import multiprocessing
# import tensorflow as tf
#
# from tensorlayerx.nn import Module
# import tensorlayerx as tlx
# from tensorlayerx.nn import (Conv2d, Linear, Flatten, MaxPool2d, BatchNorm2d)
#
# # enable debug logging
# tlx.logging.set_verbosity(tlx.logging.DEBUG)
#
# # prepare cifar10 data
# X_train, y_train, X_test, y_test = tlx.files.load_cifar10_dataset(shape=(-1, 32, 32, 3), plotable=False)
#
#
# class CNN(Module):
#
#     def __init__(self):
#         super(CNN, self).__init__()
#         # weights init
#         W_init = tlx.nn.initializers.truncated_normal(stddev=5e-2)
#         W_init2 = tlx.nn.initializers.truncated_normal(stddev=0.04)
#         b_init2 = tlx.nn.initializers.constant(value=0.1)
#
#         self.conv1 = Conv2d(64, (5, 5), (1, 1), padding='SAME', W_init=W_init, b_init=None, name='conv1', in_channels=3)
#         self.bn = BatchNorm2d(num_features=64, act=tlx.nn.ReLU)
#         self.maxpool1 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool1')
#
#         self.conv2 = Conv2d(
#             64, (5, 5), (1, 1), padding='SAME', act=tlx.nn.ReLU, W_init=W_init, b_init=None, name='conv2', in_channels=64
#         )
#         self.maxpool2 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool2')
#
#         self.flatten = Flatten(name='flatten')
#         self.linear1 = Linear(384, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear1relu', in_channels=2304)
#         self.linear2 = Linear(192, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear2relu', in_channels=384)
#         self.linear3 = Linear(10, act=None, W_init=W_init2, name='output', in_channels=192)
#
#     def forward(self, x):
#         z = self.conv1(x)
#         z = self.bn(z)
#         z = self.maxpool1(z)
#         z = self.conv2(z)
#         z = self.maxpool2(z)
#         z = self.flatten(z)
#         z = self.linear1(z)
#         z = self.linear2(z)
#         z = self.linear3(z)
#         return z
#
#
# # get the network
# net = CNN()
#
# # training settings
# batch_size = 128
# n_epoch = 500
# learning_rate = 0.0001
# print_freq = 5
# n_step_epoch = int(len(y_train) / batch_size)
# n_step = n_epoch * n_step_epoch
# shuffle_buffer_size = 128
#
# train_weights = net.trainable_weights
# optimizer = tlx.optimizers.Adam(learning_rate)
# # looking for decay learning rate? see https://github.com/tensorlayer/srgan/blob/master/train.py
#
#
# def generator_train():
#     inputs = X_train
#     targets = y_train
#     if len(inputs) != len(targets):
#         raise AssertionError("The length of inputs and targets should be equal")
#     for _input, _target in zip(inputs, targets):
#         # yield _input.encode('utf-8'), _target.encode('utf-8')
#         yield _input, _target
#
#
# def generator_test():
#     inputs = X_test
#     targets = y_test
#     if len(inputs) != len(targets):
#         raise AssertionError("The length of inputs and targets should be equal")
#     for _input, _target in zip(inputs, targets):
#         # yield _input.encode('utf-8'), _target.encode('utf-8')
#         yield _input, _target
#
#
# def _map_fn_train(img, target):
#     # 1. Randomly crop a [height, width] section of the image.
#     img = tf.image.random_crop(img, [24, 24, 3])
#     # 2. Randomly flip the image horizontally.
#     img = tf.image.random_flip_left_right(img)
#     # 3. Randomly change brightness.
#     img = tf.image.random_brightness(img, max_delta=63)
#     # 4. Randomly change contrast.
#     img = tf.image.random_contrast(img, lower=0.2, upper=1.8)
#     # 5. Subtract off the mean and divide by the variance of the pixels.
#     img = tf.image.per_image_standardization(img)
#     target = tf.reshape(target, ())
#     return img, target
#
#
# def _map_fn_test(img, target):
#     # 1. Crop the central [height, width] of the image.
#     img = tf.image.resize_with_pad(img, 24, 24)
#     # 2. Subtract off the mean and divide by the variance of the pixels.
#     img = tf.image.per_image_standardization(img)
#     img = tf.reshape(img, (24, 24, 3))
#     target = tf.reshape(target, ())
#     return img, target
#
#
# # dataset API and augmentation
# train_ds = tf.data.Dataset.from_generator(
#     generator_train, output_types=(tf.float32, tf.int32)
# )  # , output_shapes=((24, 24, 3), (1)))
# train_ds = train_ds.map(_map_fn_train, num_parallel_calls=multiprocessing.cpu_count())
# # train_ds = train_ds.repeat(n_epoch)
# train_ds = train_ds.shuffle(shuffle_buffer_size)
# train_ds = train_ds.prefetch(buffer_size=4096)
# train_ds = train_ds.batch(batch_size)
# # value = train_ds.make_one_shot_iterator().get_next()
#
# test_ds = tf.data.Dataset.from_generator(
#     generator_test, output_types=(tf.float32, tf.int32)
# )  # , output_shapes=((24, 24, 3), (1)))
# # test_ds = test_ds.shuffle(shuffle_buffer_size)
# test_ds = test_ds.map(_map_fn_test, num_parallel_calls=multiprocessing.cpu_count())
# # test_ds = test_ds.repeat(n_epoch)
# test_ds = test_ds.prefetch(buffer_size=4096)
# test_ds = test_ds.batch(batch_size)
# # value_test = test_ds.make_one_shot_iterator().get_next()
#
# for epoch in range(n_epoch):
#     start_time = time.time()
#
#     train_loss, train_acc, n_iter = 0, 0, 0
#     for X_batch, y_batch in train_ds:
#         net.set_train()
#
#         with tf.GradientTape() as tape:
#             # compute outputs
#             _logits = net(X_batch)
#             # compute loss and update model
#             _loss_ce = tlx.losses.softmax_cross_entropy_with_logits(_logits, y_batch)
#
#         grad = tape.gradient(_loss_ce, train_weights)
#         optimizer.apply_gradients(zip(grad, train_weights))
#
#         train_loss += _loss_ce
#         train_acc += np.mean(np.equal(np.argmax(_logits, 1), y_batch))
#         n_iter += 1
#
#         print("Epoch {} of {} took {}".format(epoch + 1, n_epoch, time.time() - start_time))
#         print("   train loss: {}".format(train_loss / n_iter))
#         print("   train acc:  {}".format(train_acc / n_iter))
#
#     # use training and evaluation sets to evaluate the model every print_freq epoch
#     if epoch + 1 == 1 or (epoch + 1) % print_freq == 0:
#
#         net.set_eval()
#         val_loss, val_acc, n_iter = 0, 0, 0
#         for X_batch, y_batch in test_ds:
#             _logits = net(X_batch)  # is_train=False, disable dropout
#             val_loss += tlx.losses.softmax_cross_entropy_with_logits(_logits, y_batch)
#             val_acc += np.mean(np.equal(np.argmax(_logits, 1), y_batch))
#             n_iter += 1
#         print("   val loss: {}".format(val_loss / n_iter))
#         print("   val acc:  {}".format(val_acc / n_iter))
#
# # use testing data to evaluate the model
# net.set_eval()
# test_loss, test_acc, n_iter = 0, 0, 0
# for X_batch, y_batch in test_ds:
#     _logits = net(X_batch)
#     test_loss += tlx.losses.softmax_cross_entropy_with_logits(_logits, y_batch)
#     test_acc += np.mean(np.equal(np.argmax(_logits, 1), y_batch))
#     n_iter += 1
# print("   test loss: {}".format(test_loss / n_iter))
# print("   test acc:  {}".format(test_acc / n_iter))

################################ TensorLayerX and MindSpore can be mixed programming. #################################
# import os
# os.environ['TL_BACKEND'] = 'mindspore'
#
# import time
# import numpy as np
# from tensorlayerx.nn import Module
# import tensorlayerx as tlx
# from tensorlayerx.nn import (Conv2d, Linear, Flatten, MaxPool2d, BatchNorm2d)
# from tensorlayerx.vision.transforms import (
#     Compose, Resize, RandomFlipHorizontal, RandomContrast, RandomBrightness, StandardizePerImage, RandomCrop, HWC2CHW
# )
# from tensorlayerx.dataflow import Dataset, DataLoader
# from mindspore.nn import WithLossCell, Adam
# from mindspore import ParameterTuple
# import mindspore.nn as nn
# from mindspore.ops import composite as C
# import mindspore.ops.operations as P
#
# # enable debug logging
# tlx.logging.set_verbosity(tlx.logging.DEBUG)
#
#
# class CNN(Module):
#
#     def __init__(self):
#         super(CNN, self).__init__()
#         self.conv1 = Conv2d(
#             64, (5, 5), (2, 2), b_init=None, name='conv1', in_channels=3, act=tlx.nn.ReLU, data_format='channels_first'
#         )
#         self.bn = BatchNorm2d(num_features=64, act=tlx.nn.ReLU, data_format='channels_first')
#         self.maxpool1 = MaxPool2d((3, 3), (2, 2), name='pool1', data_format='channels_first')
#         self.conv2 = Conv2d(
#             128, (5, 5), (2, 2), act=tlx.nn.ReLU, b_init=None, name='conv2', in_channels=64, data_format='channels_first'
#         )
#         self.maxpool2 = MaxPool2d((3, 3), (2, 2), name='pool2', data_format='channels_first')
#
#         self.flatten = Flatten(name='flatten')
#         self.linear1 = Linear(120, act=tlx.nn.ReLU, name='linear1relu', in_channels=512)
#         self.linear2 = Linear(84, act=tlx.nn.ReLU, name='linear2relu', in_channels=120)
#         self.linear3 = Linear(10, act=None, name='output', in_channels=84)
#
#     def forward(self, x):
#         z = self.conv1(x)
#         z = self.bn(z)
#         z = self.maxpool1(z)
#         z = self.conv2(z)
#         z = self.maxpool2(z)
#         z = self.flatten(z)
#         z = self.linear1(z)
#         z = self.linear2(z)
#         z = self.linear3(z)
#         return z
#
#
# # training settings
# batch_size = 128
# n_epoch = 500
# shuffle_buffer_size = 128
#
# # prepare cifar10 data
# X_train, y_train, X_test, y_test = tlx.files.load_cifar10_dataset(shape=(-1, 32, 32, 3), plotable=False)
#
#
# class make_dataset(Dataset):
#
#     def __init__(self, data, label, transforms):
#         self.data = data
#         self.label = label
#         self.transforms = transforms
#
#     def __getitem__(self, idx):
#         x = self.data[idx].astype('uint8')
#         y = self.label[idx].astype('int64')
#         x = self.transforms(x)
#
#         return x, y
#
#     def __len__(self):
#
#         return len(self.label)
#
#
# train_transforms = Compose(
#     [
#         RandomCrop(size=[24, 24]),
#         RandomFlipHorizontal(),
#         RandomBrightness(brightness_factor=(0.5, 1.5)),
#         RandomContrast(contrast_factor=(0.5, 1.5)),
#         StandardizePerImage(),
#         HWC2CHW()
#     ]
# )
#
# test_transforms = Compose([Resize(size=(24, 24)), StandardizePerImage(), HWC2CHW()])
#
# train_dataset = make_dataset(data=X_train, label=y_train, transforms=train_transforms)
# test_dataset = make_dataset(data=X_test, label=y_test, transforms=test_transforms)
#
# train_dataset = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# test_dataset = DataLoader(test_dataset, batch_size=batch_size)
#
#
# class GradWrap(Module):
#     """ GradWrap definition """
#
#     def __init__(self, network):
#         super(GradWrap, self).__init__(auto_prefix=False)
#         self.network = network
#         self.weights = ParameterTuple(filter(lambda x: x.requires_grad, network.get_parameters()))
#
#     def forward(self, x, label):
#         return C.GradOperation(get_by_list=True)(self.network, self.weights)(x, label)
#
#
# # get the network
# net = CNN()
# train_weights = net.trainable_weights
# optimizer = Adam(train_weights, learning_rate=0.01)
# # optimizer = Momentum(train_weights, 0.01, 0.5)
# criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
# net_with_criterion = WithLossCell(net, criterion)
# train_network = GradWrap(net_with_criterion)
# train_network.set_train()
# # print(train_weights)
# for epoch in range(n_epoch):
#     start_time = time.time()
#     train_network.set_train()
#     train_loss, train_acc, n_iter = 0, 0, 0
#     for X_batch, y_batch in train_dataset:
#         output = net(X_batch)
#         loss_output = criterion(output, y_batch)
#         grads = train_network(X_batch, y_batch)
#         success = optimizer(grads)
#         loss = loss_output.asnumpy()
#         train_loss += loss
#         n_iter += 1
#         train_acc += np.mean((P.Equal()(P.Argmax(axis=1)(output), y_batch).asnumpy()))
#         print("Epoch {} of {} took {}".format(epoch + 1, n_epoch, time.time() - start_time))
#         print("   train loss: {}".format(train_loss / n_iter))
#         print("   train acc:  {}".format(train_acc / n_iter))
#         print(" loss ", loss)

################################### TensorLayerX and Paddle can be mixed programming. ##################################
# import os
# os.environ['TL_BACKEND'] = 'paddle'

# import time
# import paddle as pd
# from tensorlayerx.nn import Module
# import tensorlayerx as tlx
# from tensorlayerx.dataflow import Dataset, DataLoader
# from tensorlayerx.nn import (Conv2d, Linear, Flatten, MaxPool2d, BatchNorm2d)
# from tensorlayerx.vision.transforms import (
#     Compose, Resize, RandomFlipHorizontal, RandomContrast, RandomBrightness, StandardizePerImage, RandomCrop
# )
# # enable debug logging
# tlx.logging.set_verbosity(tlx.logging.DEBUG)

# # prepare cifar10 data
# X_train, y_train, X_test, y_test = tlx.files.load_cifar10_dataset(shape=(-1, 32, 32, 3), plotable=False)


# class CNN(Module):

#     def __init__(self):
#         super(CNN, self).__init__()
#         # weights init
#         W_init = tlx.nn.initializers.truncated_normal(stddev=5e-2)
#         W_init2 = tlx.nn.initializers.truncated_normal(stddev=0.04)
#         b_init2 = tlx.nn.initializers.constant(value=0.1)

#         self.conv1 = Conv2d(64, (5, 5), (1, 1), padding='SAME', W_init=W_init, b_init=None, name='conv1', in_channels=3)
#         self.bn1 = BatchNorm2d(num_features=64, act=tlx.nn.ReLU)
#         self.maxpool1 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool1')

#         self.conv2 = Conv2d(
#             64, (5, 5), (1, 1), padding='SAME', W_init=W_init, b_init=None, name='conv2', in_channels=64
#         )
#         self.bn2 = BatchNorm2d(num_features=64, act=tlx.nn.ReLU)
#         self.maxpool2 = MaxPool2d((3, 3), (2, 2), padding='SAME', name='pool2')

#         self.flatten = Flatten(name='flatten')
#         self.linear1 = Linear(384, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear1relu', in_channels=2304)
#         self.linear2 = Linear(192, act=tlx.nn.ReLU, W_init=W_init2, b_init=b_init2, name='linear2relu', in_channels=384)
#         self.linear3 = Linear(10, act=None, W_init=W_init2, name='output', in_channels=192)

#     def forward(self, x):
#         z = self.conv1(x)
#         z = self.bn1(z)
#         z = self.maxpool1(z)
#         z = self.conv2(z)
#         z = self.bn2(z)
#         z = self.maxpool2(z)
#         z = self.flatten(z)
#         z = self.linear1(z)
#         z = self.linear2(z)
#         z = self.linear3(z)
#         return z


# # get the network
# net = CNN()

# # training settings
# batch_size = 128
# n_epoch = 500
# learning_rate = 0.0001
# print_freq = 5
# shuffle_buffer_size = 128
# metrics = tlx.metrics.Accuracy()

# train_weights = net.trainable_weights
# optimizer = tlx.optimizers.Adam(learning_rate)
# # looking for decay learning rate? see https://github.com/tensorlayer/srgan/blob/master/train.py


# class make_dataset(Dataset):

#     def __init__(self, data, label, transforms):
#         self.data = data
#         self.label = label
#         self.transforms = transforms

#     def __getitem__(self, idx):
#         x = self.data[idx].astype('uint8')
#         y = self.label[idx].astype('int64')
#         x = self.transforms(x)

#         return x, y

#     def __len__(self):

#         return len(self.label)


# train_transforms = Compose(
#     [
#         RandomCrop(size=[24, 24]),
#         RandomFlipHorizontal(),
#         RandomBrightness(brightness_factor=(0.5, 1.5)),
#         RandomContrast(contrast_factor=(0.5, 1.5)),
#         StandardizePerImage()
#     ]
# )

# test_transforms = Compose([Resize(size=(24, 24)), StandardizePerImage()])

# train_dataset = make_dataset(data=X_train, label=y_train, transforms=train_transforms)
# test_dataset = make_dataset(data=X_test, label=y_test, transforms=test_transforms)

# train_dataset = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# test_dataset = DataLoader(test_dataset, batch_size=batch_size)

# for epoch in range(n_epoch):
#     train_loss, train_acc, n_iter = 0, 0, 0
#     start_time = time.time()
#     for X_batch, y_batch in train_dataset:
#         net.set_train()
#         output = net(X_batch)
#         loss = pd.nn.functional.cross_entropy(output, y_batch)
#         loss_ce = loss.numpy()
#         grads = optimizer.gradient(loss, train_weights)
#         optimizer.apply_gradients(zip(grads, train_weights))

#         train_loss += loss_ce

#         if metrics:
#             metrics.update(output, y_batch)
#             train_acc += metrics.result()
#             metrics.reset()
#         n_iter += 1

#         print("Epoch {} of {} took {}".format(epoch + 1, n_epoch, time.time() - start_time))
#         print("   train loss: {}".format(train_loss / n_iter))
#         print("   train acc:  {}".format(train_acc / n_iter))