| 31 | |
| 32 | """ |
| 33 | def get_network(num_in_channels: int, args: argparse.Namespace): |
| 34 | # Define a conv net for estimating the probabilities at each decision node |
| 35 | features = base_architecture_to_features[args.net](pretrained=not args.disable_pretrained) |
| 36 | features_name = str(features).upper() |
| 37 | if features_name.startswith('VGG') or features_name.startswith('RES'): |
| 38 | first_add_on_layer_in_channels = \ |
| 39 | [i for i in features.modules() if isinstance(i, nn.Conv2d)][-1].out_channels |
| 40 | elif features_name.startswith('DENSE'): |
| 41 | first_add_on_layer_in_channels = \ |
| 42 | [i for i in features.modules() if isinstance(i, nn.BatchNorm2d)][-1].num_features |
| 43 | else: |
| 44 | raise Exception('other base base_architecture NOT implemented') |
| 45 | |
| 46 | add_on_layers = nn.Sequential( |
| 47 | nn.Conv2d(in_channels=first_add_on_layer_in_channels, out_channels=args.num_features, kernel_size=1, bias=False), |
| 48 | nn.Sigmoid() |
| 49 | ) |
| 50 | return features, add_on_layers |
| 51 | |
| 52 | def freeze(tree: ProtoTree, epoch: int, params_to_freeze: list, params_to_train: list, args: argparse.Namespace, log: Log): |
| 53 | if args.freeze_epochs>0: |