| import torch | |
| import torch.nn as nn | |
| from rscd.models.backbones.resnet import get_resnet18, get_resnet50_OS32, get_resnet50_OS8 | |
| from rscd.models.backbones.swintransformer import * | |
| class Resnet18(nn.Module): | |
| def __init__(self): | |
| super(Resnet18, self).__init__() | |
| self.backbone = get_resnet18(pretrained=True) | |
| def forward(self, xA, xB): | |
| _, xA1, xA2, xA3, xA4 = self.backbone(xA) | |
| _, xB1, xB2, xB3, xB4 = self.backbone(xB) | |
| return [xA1, xA2, xA3, xA4, xB1, xB2, xB3, xB4] | |
| class Swin(nn.Module): | |
| def __init__(self): | |
| super(Swin, self).__init__() | |
| self.backbone = swin_tiny(True) | |
| def forward(self, xA, xB): | |
| xA1, xA2, xA3, xA4 = self.backbone(xA) | |
| xB1, xB2, xB3, xB4 = self.backbone(xB) | |
| return [xA1, xA2, xA3, xA4, xB1, xB2, xB3, xB4] | |