(self, bb_pretrained=True)
| 21 | |
| 22 | class BiRefNet(nn.Module): |
| 23 | def __init__(self, bb_pretrained=True): |
| 24 | super(BiRefNet, self).__init__() |
| 25 | self.config = Config() |
| 26 | self.epoch = 1 |
| 27 | self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained) |
| 28 | |
| 29 | channels = self.config.lateral_channels_in_collection |
| 30 | |
| 31 | if self.config.auxiliary_classification: |
| 32 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) |
| 33 | self.cls_head = nn.Sequential( |
| 34 | nn.Linear(channels[0], len(class_labels_TR_sorted)) |
| 35 | ) |
| 36 | |
| 37 | if self.config.squeeze_block: |
| 38 | self.squeeze_module = nn.Sequential(*[ |
| 39 | eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0]) |
| 40 | for _ in range(eval(self.config.squeeze_block.split('_x')[1])) |
| 41 | ]) |
| 42 | |
| 43 | self.decoder = Decoder(channels) |
| 44 | |
| 45 | if self.config.ender: |
| 46 | self.dec_end = nn.Sequential( |
| 47 | nn.Conv2d(1, 16, 3, 1, 1), |
| 48 | nn.Conv2d(16, 1, 3, 1, 1), |
| 49 | nn.ReLU(inplace=True), |
| 50 | ) |
| 51 | |
| 52 | # refine patch-level segmentation |
| 53 | if self.config.refine: |
| 54 | if self.config.refine == 'itself': |
| 55 | self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') |
| 56 | else: |
| 57 | self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1')) |
| 58 | |
| 59 | if self.config.freeze_bb: |
| 60 | # Freeze the backbone... |
| 61 | print(self.named_parameters()) |
| 62 | for key, value in self.named_parameters(): |
| 63 | if 'bb.' in key and 'refiner.' not in key: |
| 64 | value.requires_grad = False |
| 65 | |
| 66 | def forward_enc(self, x): |
| 67 | if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: |
nothing calls this directly
no test coverage detected