import paddle
import paddle.nn.functional
#组网
class AlexNet(paddle.nn.Layer):
    def __init__(self, class_dim):
        super(AlexNet, self).__init__()

        self.conv1 = paddle.nn.Conv2D(3,96,11,stride=4,padding=2)
        self.pool1 = paddle.nn.layer.Pool2D(pool_size=3,pool_type='max', pool_stride=2, pool_padding=0)
        self.conv2 = paddle.nn.Conv2D(96,256,5,stride=1,padding=2)
        self.pool2 = paddle.nn.layer.Pool2D(pool_size=3,pool_type='max', pool_stride=2, pool_padding=0)
        self.conv3 = paddle.nn.Conv2D(256,384,3,stride=1,padding=1)
        self.conv4 = paddle.nn.Conv2D(384,384,3,stride=1,padding=1)
        self.conv5 = paddle.nn.Conv2D(384,256,3,stride=1,padding=1)
        self.pool3 = paddle.nn.layer.Pool2D(pool_size=3,pool_type='max', pool_stride=2, pool_padding=0)

        self.fc1 = paddle.nn.Linear(9216, 4096)
        self.fc2 = paddle.nn.Linear(4096, 4096)
        self.fc3 = paddle.nn.Linear(4096, class_dim)

        self.relu = paddle.nn.ReLU()
        self.softmax = paddle.nn.Softmax()

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.pool2(x)
        x = self.conv3(x)
        x = self.relu(x)
        x = self.conv4(x)
        x = self.relu(x)
        x = self.conv5(x)
        x = self.relu(x)
        x = self.pool3(x)
        x = paddle.flatten(x, start_axis= 1, stop_axis= 3)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        x = self.softmax(x)
        return x

class_dim=train_parameters['class_dim']
model = paddle.Model(AlexNet(class_dim=class_dim))
model.summary((None,3, 224, 224))