| import torch |
| import torch.nn as nn |
|
|
| from .subnetwork_utils import BaseBlockConvBN, TopHead |
| from .register_modules import register_model |
|
|
|
|
| class BaseSubNetwork(nn.Module): |
| def __init__(self, |
| input_channels: int, |
| base_channels: int, |
| fc_hidden_units: int, |
| fc_pred_units: int, |
| pred_activation: str, |
| ): |
| super(BaseSubNetwork, self).__init__() |
|
|
| self.input_channels = input_channels |
| self.base_channels = base_channels |
| self.fc_hidden_units = fc_hidden_units |
| self.fc_pred_units = fc_pred_units |
| self.pred_activation = pred_activation |
|
|
| self.intermediate_features = None |
|
|
| def get_intermediate_features(self) -> torch.Tensor: |
| return self.intermediate_features |
|
|
|
|
| @register_model("base_one") |
| class Subnet(BaseSubNetwork): |
| def __init__(self, input_channels=3, base_channels=32, fc_hidden_units=64, fc_pred_units=1, pred_activation="sigmoid"): |
|
|
| super(Subnet, self).__init__( |
| input_channels=input_channels, |
| base_channels=base_channels, |
| fc_hidden_units=fc_hidden_units, |
| fc_pred_units=fc_pred_units, |
| pred_activation=pred_activation, |
| ) |
|
|
| self.block_one = BaseBlockConvBN(in_ch=input_channels, |
| out_ch=base_channels, |
| conv_layers=2, |
| kernel_size=(3, 3), |
| stride=(2, 2), |
| padding=(1, 1), |
| activation="relu", |
| normalization=True,) |
|
|
| self.block_two = BaseBlockConvBN(in_ch=base_channels, |
| out_ch=base_channels*2, |
| conv_layers=2, |
| kernel_size=(3, 3), |
| stride=(2, 2), |
| padding=(1, 1), |
| activation="relu", |
| normalization=True,) |
| |
| self.block_three = BaseBlockConvBN(in_ch=base_channels*2, |
| out_ch=base_channels*4, |
| conv_layers=3, |
| kernel_size=(3, 3), |
| stride=(2, 2), |
| padding=(1, 1), |
| activation="relu", |
| normalization=True,) |
| self.flatten = nn.Flatten() |
|
|
| self.head = TopHead(fc_units=fc_hidden_units, |
| num_classes=fc_pred_units, |
| hidden_layers=1, |
| dropout_rate=0.6, |
| fc_activation="relu", |
| pred_activation=pred_activation) |
|
|
| self.intermediate_features = None |
|
|
| def forward(self, x): |
| x = self.block_one(x) |
| x = self.block_two(x) |
| self.intermediate_features = self.block_two.get_block_feats() |
| x = self.block_three(x) |
| x = self.flatten(x) |
| x = self.head(x) |
| return x |
|
|
