#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2022/5/6 09:27
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : nn_sequential.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""
import torch.nn as nn
from torch import ones
from torch.utils.tensorboard import SummaryWriter


class MyCLFNet(nn.Module):
    def __init__(self):
        super(MyCLFNet, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=(5, 5), padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(32, 32, kernel_size=(5, 5), padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(32, 64, kernel_size=(5, 5), padding=2),
            nn.MaxPool2d(kernel_size=2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )

    def forward(self, x):
        x = self.model1(x)
        return x


my_clf_net = MyCLFNet()
# 输出网络结构
print(my_clf_net)
# 快速检验网络架构正确性
input_tensor = ones((64, 3, 32, 32))
output_tensor = my_clf_net(input_tensor)
writer = SummaryWriter(log_dir="myclfnet_logs")
writer.add_graph(my_clf_net, input_tensor)
writer.close()