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

def conv_block(input_channels, num_channels):
	return nn.Sequential(
		nn.BatchNorm2d(input_channels), nn.ReLU(),
		nn.Conv2d(input_channels, num_channels, kernel_size=3, padding=1)
	)

class DenseBlock(nn.Module):
	def __init__(self, num_convs, input_channels, num_channels):
		super(Denseblock, self).__init__()
		layer = []
		for i in range(num_convs):
			layer.append(conv_block(num_channels * i + input_channels, num_channels))
		self.net = nn.Sequential(*layer)

	def forward(self, X):
		for blk in self.net:
			Y = blk(X)
			# 連結(jié)通道維度上的每個(gè)塊的輸入和輸出
			X = torch.cat((X, Y), dim=1)
		return X

blk = DenseBlock(2, 3, 10)
X = torch.randn(4, 3, 8, 8)
Y = blk(X)
Y.shape

torch.Size([4, 23, 8, 8])

def transition_block(input_channels, num_channels):
	return nn.Sequential(
		nn.BatchNorm2d(input_channels), nn.ReLU(),
		nn.Conv2d(input_channels, num_channels, kernel_size=1)
		nn.AvgPool2d(kernel_size=2, stride=2)
	)

blk = transition_block(23, 10)
blk(Y).shape

torch.Size([4, 10, 4, 4])

b1 = nn.Sequential(
	nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3),
	nn.BatchNorm2d(64), nn.ReLU(),
	nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
)

# 'num_channels'為當(dāng)前通道數(shù)
num_channels, growth_rate = 64, 32
num_convs_in_dense_blocks = [4, 4, 4, 4]
blks = []
for i, num_convs in enumerate(num_convs_in_dense_blocks):
	blks.append(DenseBlock(num_convs, num_channels, growth_rate))
	# 上一個(gè)稠密塊的輸出通道數(shù)
	num_channels += num_convs * growth_rate
	# 在稠密塊之間添加一個(gè)轉(zhuǎn)換層，使通道數(shù)量減半
	if i != len(num_convs_in_dense_blocks) - 1:
		blks.append(transition_block(num_channels, num_channels // 2))
		num_channels = num_channels // 2

net = nn.Sequential(
	b1, *blks,
	nn.BatchNorm2d(num_channels), nn.ReLU(),
	nn.AdaptiveMaxPool2d((1, 1)),
	nn.Flatten(),
	nn.Linear(num_channels, 10)
)

lr, num_epochs, batch_size = 0.1, 10, 256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size, resize=96)
d2l.train_ch6(net, train_iter, test_iter, num_epochs, lr, d2l.try_gpu())

loss 0.154, train acc 0.943, test acc 0.880
5506.9 examples/sec on cuda:0