(self, args)
| 26 | |
| 27 | class BiRAFT(nn.Module): |
| 28 | def __init__(self, args): |
| 29 | super(BiRAFT, self).__init__() |
| 30 | self.args = args |
| 31 | |
| 32 | if args.small: |
| 33 | self.hidden_dim = hdim = 96 |
| 34 | self.context_dim = cdim = 64 |
| 35 | args.corr_levels = 4 |
| 36 | args.corr_radius = 3 |
| 37 | |
| 38 | else: |
| 39 | self.hidden_dim = hdim = 128 |
| 40 | self.context_dim = cdim = 128 |
| 41 | args.corr_levels = 4 |
| 42 | args.corr_radius = 4 |
| 43 | |
| 44 | if 'dropout' not in self.args: |
| 45 | self.args.dropout = 0 |
| 46 | |
| 47 | if 'alternate_corr' not in self.args: |
| 48 | self.args.alternate_corr = False |
| 49 | |
| 50 | # feature network, context network, and update block |
| 51 | if args.small: |
| 52 | self.fnet = SmallEncoder(output_dim=128, norm_fn='instance', dropout=args.dropout) |
| 53 | self.cnet = SmallEncoder(output_dim=hdim + cdim, norm_fn='none', dropout=args.dropout) |
| 54 | self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim) |
| 55 | |
| 56 | else: |
| 57 | if self.args.fnet == 'CNN': |
| 58 | self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=args.dropout) |
| 59 | self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='batch', dropout=args.dropout) |
| 60 | elif self.args.fnet == 'twins': |
| 61 | self.fnet = twins_svt_large(pretrained=True) |
| 62 | self.cnet = twins_svt_large(pretrained=True) |
| 63 | self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim) |
| 64 | |
| 65 | def freeze_bn(self): |
| 66 | for m in self.modules(): |
nothing calls this directly
no test coverage detected