File size: 3,245 Bytes
46fdf2a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
from torchvision.models import ResNet
from torchvision.models.resnet import Bottleneck, BasicBlock
from .csra import CSRA, MHA
import torch.utils.model_zoo as model_zoo
import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
model_urls = {
"resnet18": "https://download.pytorch.org/models/resnet18-5c106cde.pth",
"resnet34": "https://download.pytorch.org/models/resnet34-333f7ec4.pth",
"resnet50": "https://download.pytorch.org/models/resnet50-19c8e357.pth",
"resnet101": "https://download.pytorch.org/models/resnet101-5d3b4d8f.pth",
"resnet152": "https://download.pytorch.org/models/resnet152-b121ed2d.pth",
}
class ResNet_CSRA(ResNet):
arch_settings = {
18: (BasicBlock, (2, 2, 2, 2)),
34: (BasicBlock, (3, 4, 6, 3)),
50: (Bottleneck, (3, 4, 6, 3)),
101: (Bottleneck, (3, 4, 23, 3)),
152: (Bottleneck, (3, 8, 36, 3)),
}
def __init__(
self, num_heads, lam, num_classes, depth=101, input_dim=2048, cutmix=None
):
self.block, self.layers = self.arch_settings[depth]
self.depth = depth
super(ResNet_CSRA, self).__init__(self.block, self.layers)
self.init_weights(pretrained=True, cutmix=cutmix)
self.classifier = MHA(num_heads, lam, input_dim, num_classes)
self.loss_func = F.binary_cross_entropy_with_logits
# todo
# criterion = nn.BCEWithLogitsLoss() # loss combines a Sigmoid layer and the BCELoss in one single class
def backbone(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
return x
def forward_train(self, x, target):
x = self.backbone(x)
logit = self.classifier(x)
loss = self.loss_func(logit, target, reduction="mean")
return logit, loss
def forward_test(self, x):
x = self.backbone(x)
x = self.classifier(x)
return x
def forward(self, x, target=None):
if target is not None:
return self.forward_train(x, target)
else:
return self.forward_test(x)
def init_weights(self, pretrained=True, cutmix=None):
if cutmix is not None:
print("backbone params inited by CutMix pretrained model")
state_dict = torch.load(cutmix)
elif pretrained:
print("backbone params inited by Pytorch official model")
model_url = model_urls["resnet{}".format(self.depth)]
state_dict = model_zoo.load_url(model_url)
model_dict = self.state_dict()
try:
pretrained_dict = {k: v for k, v in state_dict.items() if k in model_dict}
self.load_state_dict(pretrained_dict)
except:
logger = logging.getLogger()
logger.info(
"the keys in pretrained model is not equal to the keys in the ResNet you choose, trying to fix..."
)
state_dict = self._keysFix(model_dict, state_dict)
self.load_state_dict(state_dict)
# remove the original 1000-class fc
self.fc = nn.Sequential()
|