| import torch | |
| from torch import nn | |
| from architectures.SLDDLevel import SLDDLevel | |
| class FinalLayer(): | |
| def __init__(self, num_classes, n_features): | |
| super().__init__() | |
| self.avgpool = torch.nn.AdaptiveAvgPool2d((1, 1)) | |
| self.linear = nn.Linear(n_features, num_classes) | |
| self.featureDropout = torch.nn.Dropout(0.2) | |
| self.selection = None | |
| def transform_output(self, feature_maps, with_feature_maps, | |
| with_final_features): | |
| if self.selection is not None: | |
| feature_maps = feature_maps[:, self.selection] | |
| x = self.avgpool(feature_maps) | |
| pre_out = torch.flatten(x, 1) | |
| final_features = self.featureDropout(pre_out) | |
| final = self.linear(final_features) | |
| final = [final] | |
| if with_feature_maps: | |
| final.append(feature_maps) | |
| if with_final_features: | |
| final.append(final_features) | |
| if len(final) == 1: | |
| final = final[0] | |
| return final | |
| def set_model_sldd(self, selection, weight, mean, std, bias = None): | |
| self.selection = selection | |
| self.linear = SLDDLevel(selection, weight, mean, std, bias) | |
| self.featureDropout = torch.nn.Dropout(0.1) |