| import torch |
| import torch.nn as nn |
| from torch import Tensor |
|
|
|
|
| class Fire(nn.Module): |
|
|
| def __init__(self, inplanes: int, squeeze_planes: int, expand1x1_planes: int, expand3x3_planes: int): |
| super(Fire, self).__init__() |
| self.inplanes = inplanes |
| self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1) |
| self.squeeze_activation = nn.ReLU(inplace=True) |
| self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes, kernel_size=1) |
| self.expand1x1_activation = nn.ReLU(inplace=True) |
| self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes, kernel_size=3, padding=1) |
| self.expand3x3_activation = nn.ReLU(inplace=True) |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| x = self.squeeze_activation(self.squeeze(x)) |
| return torch.cat([self.expand1x1_activation(self.expand1x1(x)), |
| self.expand3x3_activation(self.expand3x3(x))], 1) |
|
|