DenseNet详解链接 代码： DenseNet.py
import torch
import torch.nn as nn
from collections import OrderedDict
class _DenseLayer(nn.Sequential):
def __init__(self, in_channels, growth_rate, bn_size):
super(_DenseLayer, self).__init__()
self.add_module('norm1', nn.BatchNorm2d(in_channels))
self.add_module('relu1', nn.ReLU(inplace=True))
self.add_module('conv1', nn.Conv2d(in_channels, bn_size * growth_rate,
kernel_size=1,
stride=1, bias=False))
self.add_module('norm2', nn.BatchNorm2d(bn_size*growth_rate))
self.add_module('relu2', nn.ReLU(inplace=True))
self.add_module('conv2', nn.Conv2d(bn_size*growth_rate, growth_rate,
kernel_size=3,
stride=1, padding=1, bias=False))
# 重载forward函数
def forward(self, x):
new_features = super(_DenseLayer, self).forward(x)
return torch.cat([x, new_features], 1)
class _DenseBlock(nn.Sequential):
def __init__(self, num_layers, in_channels, bn_size, growth_rate):
super(_DenseBlock, self).__init__()
for i in range(num_layers):
self.add_module('denselayer%d' % (i+1),
_DenseLayer(in_channels+growth_rate*i,
growth_rate, bn_size))
class _Transition(nn.Sequential):
def __init__(self, in_channels, out_channels):
super(_Transition, self).__init__()
self.add_module('norm', nn.BatchNorm2d(in_channels))
self.add_module('relu', nn.ReLU(inplace=True))
self.add_module('conv', nn.Conv2d(in_channels, out_channels,
kernel_size=1,
stride=1, bias=False))
self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))
class DenseNet_BC(nn.Module):
def __init__(self, growth_rate=12, block_config=(6,12,24,16),
bn_size=4, theta=0.5, num_classes=10):
super(DenseNet_BC, self).__init__()
# 初始的卷积为filter:2倍的growth_rate
num_init_feature = 2 * growth_rate
# 表示cifar-10
if num_classes == 10:
self.features = nn.Sequential(OrderedDict([
('conv0', nn.Conv2d(3, num_init_feature,
kernel_size=3, stride=1,
padding=1, bias=False)),
]))
else:
self.features = nn.Sequential(OrderedDict([
('conv0', nn.Conv2d(3, num_init_feature,
kernel_size=7, stride=2,
padding=3, bias=False)),
('norm0', nn.BatchNorm2d(num_init_feature)),
('relu0', nn.ReLU(inplace=True)),
('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
]))
num_feature = num_init_feature
for i, num_layers in enumerate(block_config):
self.features.add_module('denseblock%d' % (i+1),
_DenseBlock(num_layers, num_feature,
bn_size, growth_rate))
num_feature = num_feature + growth_rate * num_layers
if i != len(block_config)-1:
self.features.add_module('transition%d' % (i + 1),
_Transition(num_feature,
int(num_feature * theta)))
num_feature = int(num_feature * theta)
self.features.add_module('norm5', nn.BatchNorm2d(num_feature))
self.features.add_module('relu5', nn.ReLU(inplace=True))
self.features.add_module('avg_pool', nn.AdaptiveAvgPool2d((1, 1)))
self.classifier = nn.Linear(num_feature, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight)
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.constant_(m.bias, 0)
def forward(self, x):
features = self.features(x)
out = features.view(features.size(0), -1)
out = self.classifier(out)
return out
# DenseNet_BC for ImageNet
def DenseNet121():
return DenseNet_BC(growth_rate=32, block_config=(6, 12, 24, 16), num_classes=1000)
def DenseNet169():
return DenseNet_BC(growth_rate=32, block_config=(6, 12, 32, 32), num_classes=1000)
def DenseNet201():
return DenseNet_BC(growth_rate=32, block_config=(6, 12, 48, 32), num_classes=1000)
def DenseNet161():
return DenseNet_BC(growth_rate=48, block_config=(6, 12, 36, 24), num_classes=1000,)
# DenseNet_BC for cifar
def densenet_BC_100():
return DenseNet_BC(growth_rate=12, block_config=(16, 16, 16))
main.py
import torch
from torch import nn, optim
import torchvision.transforms as transforms
from torchvision import datasets
from torch.utils.data import DataLoader
from DenseNet import densenet_BC_100
import cv2
# 用CIFAR-10 数据集进行实验
def main():
batchsz = 64
cifar_train = datasets.CIFAR10('cifar', True, transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]), download=True)
cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)
cifar_test = datasets.CIFAR10('cifar', False, transform=transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]), download=True)
cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True)
device = torch.device('cuda')
model = densenet_BC_100().to(device)
criteon = nn.CrossEntropyLoss().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(100):
model.train()
for batchidx, (x, label) in enumerate(cifar_train):
x, label = x.to(device), label.to(device)
logits = model(x)
loss = criteon(logits, label)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(epoch, 'loss:', loss.item())
model.eval()
with torch.no_grad():
total_correct = 0
total_num = 0
for x, label in cifar_test:
x, label = x.to(device), label.to(device)
logits = model(x)
pred = logits.argmax(dim=1)
correct = torch.eq(pred, label).float().sum().item()
total_correct += correct
total_num += x.size(0)
acc = total_correct / total_num
print(epoch, 'test acc:', acc)
if __name__ == '__main__':
main()
