| 122 | return net, None, delta_flow |
| 123 | |
| 124 | class BasicUpdateBlock(nn.Module): |
| 125 | def __init__(self, args, hidden_dim=128, input_dim=128): |
| 126 | super(BasicUpdateBlock, self).__init__() |
| 127 | self.args = args |
| 128 | self.encoder = BasicMotionEncoder(args) |
| 129 | self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128+hidden_dim) |
| 130 | if self.args.fnet == 'CNN': |
| 131 | self.flow_head = FlowHead_cnn(hidden_dim, hidden_dim=256) |
| 132 | elif self.args.fnet == 'twins': |
| 133 | self.flow_head = FlowHead_twins(hidden_dim, hidden_dim=256) |
| 134 | |
| 135 | self.mask = nn.Sequential( |
| 136 | nn.Conv2d(128, 256, 3, padding=1), |
| 137 | nn.ReLU(inplace=True), |
| 138 | nn.Conv2d(256, 64*9, 1, padding=0)) |
| 139 | |
| 140 | def forward(self, net, inp, corr, flow, upsample=True): |
| 141 | motion_features = self.encoder(flow, corr) |
| 142 | inp = torch.cat([inp, motion_features], dim=1) |
| 143 | |
| 144 | net = self.gru(net, inp) |
| 145 | delta_flow = self.flow_head(net) |
| 146 | delta_flow = delta_flow[:, :2] |
| 147 | |
| 148 | # scale mask to balence gradients |
| 149 | mask = .25 * self.mask(net) |
| 150 | return net, mask, delta_flow |
| 151 | |
| 152 | |