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

'''
@author:ccc-ju
@email:ccc-ju@outlook.com
@software:PyCharm
@time:2020/7/10 16:40
'''
import torch
from torch import nn

class VGG(nn.Module):
  def __init__(self, num_classes=1000):
    super(VGG, self).__init__()
    layers = []
    in_dim = 3
    out_dim = 64
    # 利用循环构造卷积层，一共13个卷积层
    for i in range(13):
      layers += [nn.Conv2d(in_dim, out_dim, 3, stride=1, padding=1), nn.ReLU(inplace=True)]
      in_dim = out_dim
      # 在第2，4，7，10，13个卷积层后增加池化层
      if i == 1 or i == 3 or i == 6 or i == 9  or i == 12:
        layers += [nn.MaxPool2d(2, 2)]
        # 第10个卷积后保持和前边的通道数一致，都为512，其余加倍
        if i != 9:
          out_dim *= 2
    self.features = nn.Sequential(*layers)
    # VGGNet的3个全连接层，中间有ReLu与Dropout层
    self.classifier = nn.Sequential(
        nn.Linear(512 * 7 * 7, 4096),
        nn.ReLU(True),
        nn.Dropout(),
        nn.Linear(4096, 4096),
        nn.Dropout(),
        nn.Linear(4096, num_classes)
    )
  def forward(self, x):
    x = self.features(x)
    # 这里是将特征图的维度从[1, 512, 7, 7]变到[1, 512 * 7 * 7]
    x = x.view(x.size(0), -1)
    x = self.classifier(x)
    return x

vgg = VGG(21).cuda()
input = torch.randn(1, 3, 224, 224).cuda()
scores = vgg(input)
print(scores)
