def get_layer_groups(self, do_fc=False):
if do_fc:
return [self.fc_model]
idxs = [self.lr_cut]
c = children(self.top_model)
if len(c) == 3: c = children(c[0]) + c[1:]
lgs = list(split_by_idxs(c, idxs))
return lgs + [self.fc_model]
def __init__(self, block_wgts:[float]=[0.2,0.7,0.1], multiplier:float=1.0):
super().__init__()
m_vgg = vgg16(True)
blocks = [i-1 for i,o in enumerate(children(m_vgg)) if isinstance(o,nn.MaxPool2d)]
blocks, [m_vgg[i] for i in blocks]
layer_ids = blocks[:3]
vgg_layers = children(m_vgg)[:23]
m_vgg = nn.Sequential(*vgg_layers).cuda().eval()
set_trainable(m_vgg, False)
self.m,self.wgts = m_vgg,block_wgts
self.sfs = [SaveFeatures(m_vgg[i]) for i in layer_ids]
self.multiplier = multiplier