def __init__(self,
f,
c,
is_multi,
is_reg,
ps=None,
xtra_fc=None,
xtra_cut=0,
custom_head=None,
pretrained=True):
self.f, self.c, self.is_multi, self.is_reg, self.xtra_cut = f, c, is_multi, is_reg, xtra_cut
if xtra_fc is None: xtra_fc = [512]
if ps is None: ps = [0.25] * len(xtra_fc) + [0.5]
self.ps, self.xtra_fc = ps, xtra_fc
if f in model_meta:
cut, self.lr_cut = model_meta[f]
else:
cut, self.lr_cut = 0, 0
cut -= xtra_cut
layers = cut_model(f(pretrained), cut)
# replace first convolutional layer by 4->64 while keeping corresponding weights
# and initializing new weights with zeros
w = layers[0].weight
layers[0] = nn.Conv2d(4,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
layers[0].weight = torch.nn.Parameter(
torch.cat((w, torch.zeros(64, 1, 7, 7)), dim=1))
self.nf = model_features[f] if f in model_features else (
num_features(layers) * 2)
if not custom_head: layers += [AdaptiveConcatPool2d(), Flatten()]
self.top_model = nn.Sequential(*layers)
n_fc = len(self.xtra_fc) + 1
if not isinstance(self.ps, list): self.ps = [self.ps] * n_fc
if custom_head:
fc_layers = [custom_head]
else:
fc_layers = self.get_fc_layers()
self.n_fc = len(fc_layers)
self.fc_model = to_gpu(nn.Sequential(*fc_layers))
if not custom_head: apply_init(self.fc_model, kaiming_normal)
self.model = to_gpu(nn.Sequential(*(layers + fc_layers)))
def _get_pretrained_resnet_base(layers_cut:int=0):
f = resnet34
cut,lr_cut = model_meta[f]
cut-=layers_cut
layers = cut_model(f(True), cut)
return nn.Sequential(*layers), lr_cut