| 96 | return self.conv(x) |
| 97 | |
| 98 | class DenseLayer(nn.Module): |
| 99 | def __init__(self, in_channels, growth_rate, drop_rate=0.0): |
| 100 | super().__init__() |
| 101 | self.bn1 = nn.BatchNorm2d(in_channels) |
| 102 | self.relu1 = nn.ReLU(inplace=True) |
| 103 | self.conv1 = nn.Conv2d(in_channels, growth_rate * 4, kernel_size=1, bias=False) |
| 104 | |
| 105 | self.bn2 = nn.BatchNorm2d(growth_rate * 4) |
| 106 | self.relu2 = nn.ReLU(inplace=True) |
| 107 | self.conv2 = nn.Conv2d(growth_rate * 4, growth_rate, kernel_size=3, padding=1, bias=False) |
| 108 | |
| 109 | self.dropout = nn.Dropout2d(drop_rate) if drop_rate > 0 else nn.Identity() |
| 110 | |
| 111 | def forward(self, x): |
| 112 | out = self.conv1(self.relu1(self.bn1(x))) |
| 113 | out = self.conv2(self.relu2(self.bn2(out))) |
| 114 | out = self.dropout(out) |
| 115 | return torch.cat([x, out], 1) |
| 116 | |
| 117 | class DenseBlock(nn.Module): |
| 118 | def __init__(self, num_layers, in_channels, growth_rate, drop_rate=0.0): |