class CustomBlock_Agg(nn.Module):
def __init__(self, args):
'''
Given some a patch model, add add some FC layers and a shortcut to make whole image prediction
'''
super(CustomBlock_Agg, self).__init__()
self.args = args
if not args.use_precomputed_hiddens:
self.feat_extractor = load_model(args.patch_snapshot, args, False)
agg_layers = get_layers(args.block_layout)
self._model = ResNet(agg_layers, args)
def forward(self, x, risk_factors=None):
'''
param x: a batch of image tensors, in the order of:
returns hidden: last hidden layer of model
'''
x = x.data
if not self.args.use_precomputed_hiddens:
_, _, x = self.feat_extractor(x, risk_factors)
x = x.data
logit, hidden, x = self._model(x, risk_factors)
return logit, hidden, x
def cuda(self, device=None):
self.feat_extractor = self.feat_extractor.cuda(device)
self._model = self._model.cuda(device)
return self
class CustomResnet(nn.Module):
def __init__(self, args):
super(CustomResnet, self).__init__()
layers = get_layers(args.block_layout)
self._model = ResNet(layers, args)
model_name = args.pretrained_imagenet_model_name
if args.pretrained_on_imagenet:
load_pretrained_weights(self._model,
load_pretrained_model(model_name))
def forward(self, x, risk_factors=None, batch=None):
return self._model(x, risk_factors=risk_factors, batch=None)
def cuda(self, device=None):
self._model = self._model.cuda(device)
return self
class Default_Resnet50(nn.Module):
def __init__(self, args):
super(Default_Resnet50, self).__init__()
block_layout = [[('Bottleneck', 3)], [('Bottleneck', 4)],
[('Bottleneck', 6)], [('Bottleneck', 3)]]
layers = get_layers(block_layout)
self._model = ResNet(layers, args)
if args.pretrained_on_imagenet:
load_pretrained_weights(self._model,
load_pretrained_model('resnet50'))
def forward(self, x, risk_factors=None, batch=None):
return self._model(x, risk_factors=risk_factors, batch=batch)
def cuda(self, device=None):
self._model = self._model.cuda(device)
return self
class Default_8StageResnet36(nn.Module):
def __init__(self, args):
super(Default_8StageResnet36, self).__init__()
block_layout = [[('BasicBlock', 2)], [('BasicBlock', 2)],
[('BasicBlock', 2)], [('BasicBlock', 2)],
[('BasicBlock', 2)], [('BasicBlock', 2)],
[('BasicBlock', 2)], [('BasicBlock', 2)]]
layers = get_layers(block_layout)
self._model = ResNet(layers, args)
if args.pretrained_on_imagenet:
load_pretrained_weights(self._model,
load_pretrained_model('resnet18'))
def forward(self, x, risk_factors=None):
return self._model(x, risk_factors=risk_factors)
def cuda(self, device=None):
self._model = self._model.cuda(device)
return self