import torch
from torch import nn
from d2l import torch as d2l

net=nn.Sequential(
      nn.Conv2d(1,96,kernel_size=11,stride=4,padding=1),nn.ReLU(),
      nn.MaxPool2d(kernel_size=3,stride=2),
      nn.Conv2d(96,128*2,kernel_size=5,padding=2),nn.ReLU(),
      nn.MaxPool2d(kernel_size=3,stride=2),
      nn.Conv2d(128*2,192*2,kernel_size=3,padding=1),nn.ReLU(),
      nn.Conv2d(192*2,192*2,kernel_size=3,padding=1),nn.ReLU(),
      nn.Conv2d(192*2,128*2,kernel_size=3,padding=1),nn.ReLU(),
      nn.MaxPool2d(kernel_size=3,stride=2), # 6*6*256
      nn.Flatten(),
      nn.Linear(6400,2048*2),nn.ReLU(),nn.Dropout(p=0.5),
      nn.Linear(2048*2,2048*2),nn.ReLU(),nn.Dropout(p=0.5),
      nn.Linear(2048*2,10),nn.ReLU(),
    )


# 看每个层输出得形状
X=torch.randn(1,1,224,224)
for layer in net:
      X=layer(X)
      print(layer.__class__.__name__,'Output  shape:\t',X.shape)


batch_size=32
train_iter,test_iter=d2l.load_data_fashion_mnist(batch_size=batch_size,resize=224)

lr,num_epochs=0.01,5
d2l.train_ch6(net,train_iter,test_iter,num_epochs,lr,d2l.try_gpu())