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