| 348 | results=torch.cat((self.cosR_small(x),self.cosR_huge(x)),axis=1) |
| 349 | return results |
| 350 | class CostRegNet_refine(nn.Module): |
| 351 | def __init__(self, in_channels, base_channels,stage=0): |
| 352 | super(CostRegNet_refine, self).__init__() |
| 353 | self.cosR_small=CostRegNet_part_refine(in_channels, base_channels,stage=0) |
| 354 | self.cosR_huge=CostRegNet_part_refine(in_channels, base_channels,stage=0) |
| 355 | def forward(self, x): |
| 356 | results=torch.cat((self.cosR_small(x),self.cosR_huge(x)),axis=1) |
| 357 | return results |
| 358 | class CostRegNet_part(nn.Module): |
| 359 | def __init__(self, in_channels, base_channels,stage=0): |
| 360 | super(CostRegNet_part, self).__init__() |