import paddle
import numpy as np
import matplotlib.pyplot as plt
import paddle.nn as nn
from paddle.vision.datasets import Cifar100
from paddle.vision.transforms import Normalize
paddle.__version__

# 数据预处理定义
normalize = Normalize(mean=[0.5, 0.5, 0.5],
                        std=[0.5, 0.5, 0.5],
                        data_format='HWC')
# 训练集和验证集定义
train_cifar100 = Cifar100(mode='train', transform=normalize)
test_cifar100 = Cifar100(mode='test', transform=normalize)

'''
自定义数据集
'''
from paddle.io import Dataset
class MyDataset(paddle.io.Dataset):
    """
    步骤一：继承paddle.io.Dataset类
    """
    def __init__(self, data_):
        """
        步骤二：实现构造函数，定义数据集大小
        """
        super(MyDataset, self).__init__()
        self.data = []
        self.label = []
        for i in range(len(data_)):
            x = data_[i][0].transpose((2, 0, 1))
            y = data_[i][1]
            self.data.append(x)
            self.label.append(np.array(y).astype('int64'))
    def __getitem__(self, index):
        """
        步骤三：实现__getitem__方法，定义指定index时如何获取数据，并返回单条数据（训练数据，对应的标签）
        """
        # 返回单一数据和标签
        data = self.data[index]
        label = self.label[index]
        # 注：返回标签数据时必须是int64
        return data, np.array(label, dtype='int64')
    def __len__(self):
        """
        步骤四：实现__len__方法，返回数据集总数目
        """
        # 返回数据总数
        return len(self.data)
# 测试定义的数据集
train_dataset = MyDataset(data_=train_cifar100)
eval_dataset = MyDataset(data_=test_cifar100)
print('=============train_dataset =============')
# 输出数据集的形状和标签
print(train_dataset.__getitem__(1)[0].shape,train_dataset.__getitem__(1)[1])
# 输出数据集的长度
print(train_dataset.__len__())
print('=============eval_dataset =============')
# 输出数据集的形状和标签
for data, label in eval_dataset:
    print(data.shape, label)
    break
# 输出数据集的长度
print(eval_dataset.__len__())

# LeNet-5改
network = nn.Sequential(
    nn.Conv2D(in_channels=3, out_channels=6, kernel_size=5, stride=1, padding=0),  # C1 卷积层
    nn.Tanh(),
    nn.AvgPool2D(kernel_size=2, stride=2),  # S2 平局池化层
    nn.Sigmoid(),   # Sigmoid激活函数
    nn.Conv2D(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0),  # C3 卷积层
    nn.Tanh(),
    nn.AvgPool2D(kernel_size=2, stride=2),  # S4 平均池化层
    nn.Sigmoid(),  # Sigmoid激活函数
    nn.Conv2D(in_channels=16, out_channels=120, kernel_size=5, stride=1, padding=0), # C5 卷积层
    nn.Tanh(),
    nn.Flatten(),
    nn.Linear(in_features=120, out_features=100), # F6 全连接层
    nn.Tanh(),
    nn.Linear(in_features=100, out_features=100) # OUTPUT 全连接层
)
paddle.summary(network, (1, 3, 32, 32))

class LeNet_G(nn.Layer):
    """
    继承paddle.nn.Layer定义网络结构
    """
    def __init__(self, num_classes=100):
        """
        初始化函数
        """
        super(LeNet_G, self).__init__()
        self.conv2D1 = nn.Conv2D(in_channels=3, out_channels=6, kernel_size=5, stride=1, padding=0)  # C1 卷积层
        self.conv2D2 = nn.Conv2D(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0)  # C3 卷积层
        self.conv2D3 = nn.Conv2D(in_channels=16, out_channels=120, kernel_size=5, stride=1, padding=0)  # C5 卷积层
        self.avgpool2D1 = nn.AvgPool2D(kernel_size=2, stride=2)  # S2平均池化
        self.avgpool2D2 = nn.AvgPool2D(kernel_size=2, stride=2)  # S4 平均池化层
        self.sigmoid = nn.Sigmoid()  # 激活函数
        self.tanh = nn.Tanh()
        self.linear1 = nn.Linear(in_features=120, out_features=100)  # F6 全连接层
        self.linear2 = nn.Linear(in_features=100, out_features=num_classes) # OUTPUT 全连接层
    def forward(self, inputs):
        """
        前向计算
        """
        inputs = paddle.to_tensor(inputs)
        x = self.conv2D1(inputs)
        x = self.tanh(x)
        x = self.avgpool2D1(x)
        x = self.sigmoid(x)
        x = self.conv2D2(x)
        x = self.tanh(x)
        x = self.avgpool2D2(x)
        x = self.sigmoid(x)
        x = self.conv2D3(x)
        x = self.tanh(x)
        x = paddle.flatten(x, 1)
        x = self.linear1(x)
        x = self.tanh(x)
        y = self.linear2(x)
        return x
network_1 = LeNet_G()
paddle.summary(network_1, (1, 3, 32, 32))

# 用Model封装模型
model = paddle.Model(network_1)
# 定义损失函数
model.prepare(paddle.optimizer.Adam(learning_rate=0.0001, parameters=model.parameters()), 
              paddle.nn.CrossEntropyLoss(), 
              paddle.metric.Accuracy())
# 训练可视化VisualDL工具的回调函数
visualdl = paddle.callbacks.VisualDL(log_dir='visualdl_log')
# 启动模型全流程训练
model.fit(train_dataset,            # 训练数据集
          eval_dataset,            # 评估数据集
          epochs=50,            # 总的训练轮次
          batch_size = 256,    # 批次计算的样本量大小
          shuffle=True,             # 是否打乱样本集
          verbose=1,                # 日志展示格式
          save_dir='./CNN_points/', # 分阶段的训练模型存储路径
          callbacks=[visualdl])     # 回调函数使用
# 保存模型
  shuffle=True,             # 是否打乱样本集
          verbose=1,                # 日志展示格式
          save_dir='./CNN_points/', # 分阶段的训练模型存储路径
          callbacks=[visualdl])     # 回调函数使用
# 保存模型
model.save('cnn_model_save_dir')