File size: 738 Bytes
89fbe02 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import torch.nn as nn
class Layer(nn.Module):
def __init__(self, in_channels, out_channels):
"""
DenseNet layer with Batch Normalization, ELU activation,
Convolution, and Dropout.
Args:
in_channels (int): Number of input channels.
out_channels (int): Number of output channels. This is the growth rate.
"""
super(Layer, self).__init__()
self.block = nn.Sequential(
nn.BatchNorm2d(in_channels),
nn.ELU(inplace=True), # Exponential ReLU
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
nn.Dropout2d(p=0.2)
)
def forward(self, x):
x = self.block(x)
return x
|