densenet网络结构pytorch实现

DenseNet是一种深度卷积神经网络，通过密集连接（dense connection）的方式将不同层的特征图进行连接，以提高特征的重用和信息流动性。下面是使用PyTorch实现DenseNet网络结构的示例代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义DenseNet的基本单元
class BasicBlock(nn.Module):
    def __init__(self, in_channels, growth_rate):
        super(BasicBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_channels, growth_rate, kernel_size=3, padding=1, bias=False)
        
    def forward(self, x):
        out = self.conv1(self.relu(self.bn1(x)))
        out = torch.cat([x, out], 1)  # 在通道维度上进行连接
        return out
    
# 定义DenseNet网络结构
class DenseNet(nn.Module):
    def __init__(self, num_blocks, growth_rate=32, num_classes=10):
        super(DenseNet, self).__init__()
        self.in_channels = 64
        
        # 第一层卷积
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False)
        
        # DenseBlock层
        self.dense_blocks = nn.ModuleList([])
        self.transition_blocks = nn.ModuleList([])
        
        for i in range(len(num_blocks)):
            num_channels = self.in_channels + i * growth_rate
            self.dense_blocks.append(self._make_dense_block(num_channels, growth_rate, num_blocks[i]))
            self.in_channels = num_channels + num_blocks[i] * growth_rate
            
            if i != len(num_blocks) - 1:
                self.transition_blocks.append(self._make_transition_block(self.in_channels))
        
        self.bn = nn.BatchNorm2d(self.in_channels)
        self.relu = nn.ReLU(inplace=True)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(self.in_channels, num_classes)
        
    def forward(self, x):
        out = self.conv1(x)
        
        for i in range(len(self.dense_blocks)):
            out = self.dense_blocks[i](out)
            
            if i != len(self.dense_blocks) - 1:
                out = self.transition_blocks[i](out)
        
        out = self.relu(self.bn(out))
        out = self.avgpool(out)
        out = torch.flatten(out, 1)
        out = self.fc(out)
        
        return out
    
    def _make_dense_block(self, in_channels, growth_rate, num_blocks):
        layers = []
        for _ in range(num_blocks):
            layers.append(BasicBlock(in_channels, growth_rate))
            in_channels += growth_rate
        return nn.Sequential(*layers)
    
    def _make_transition_block(self, in_channels):
        block = nn.Sequential(
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels, in_channels // 2, kernel_size=1, bias=False),
            nn.AvgPool2d(kernel_size=2, stride=2)
        )
        return block

# 创建DenseNet模型
model = DenseNet([6, 12, 24, 16], growth_rate=32, num_classes=1000)

# 输出模型结构
print(model)