| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torchvision | |
| from rscd.models.backbones.Decom_Backbone import ResNet3D | |
| class AFCD3D_backbone(nn.Module): | |
| def __init__(self): | |
| super(AFCD3D_backbone, self).__init__() | |
| resnet = torchvision.models.resnet18(pretrained=True) | |
| self.resnet = ResNet3D(resnet) | |
| def forward(self, imageA, imageB): | |
| imageA = imageA.unsqueeze(2) | |
| imageB = imageB.unsqueeze(2) | |
| x = torch.cat([imageA, imageB], 2) | |
| size = x.size()[3:] | |
| x = self.resnet.conv1(x) | |
| x = self.resnet.bn1(x) | |
| x0 = self.resnet.relu(x) | |
| x = self.resnet.maxpool(x0) | |
| x1 = self.resnet.layer1(x) | |
| x2 = self.resnet.layer2(x1) | |
| x3 = self.resnet.layer3(x2) | |
| x4 = self.resnet.layer4(x3) | |
| return [size, x0, x1, x2, x3, x4] | |