def arch_summary(arch):
model = arch(False)
tot = 0
for i, l in enumerate(model.children()):
n_layers = len(flatten_model(l))
tot += n_layers
print(f'({i}) {l.__class__.__name__:<12}: {n_layers:<4}layers (total: {tot})')
def __init__(
self,
learn: Learner,
do_remove: bool = True,
hMin=-1,
hMax=1,
nBins=100,
useClasses=False, # if true compute histogram of classes in the last layer
liveChart=True, # show live chart of last layer
modulesId=-1, # array of modules to keep
):
self.hMin = hMin or (-hMax)
self.hMax = hMax
self.nBins = nBins
self.zero_bin = math.floor(-self.nBins * self.hMin /
(self.hMax - self.hMin))
self.liveChart = liveChart
self.allModules = [m for m in flatten_model(learn.model)]
self.useClasses = useClasses
modules = self.allModules
if modulesId:
modules = [self.allModules[i] for i in listify(modulesId)]
self.allModules = modules if modules else self.allModules
self.c = learn.data.c # Number of Calsses
super().__init__(learn, modules, do_remove)
def _freeze(self):
"Freezes the pretrained backbone."
for idx, i in enumerate(flatten_model(self.learn.model.backbone)):
if isinstance(i, (torch.nn.BatchNorm2d)):
continue
for p in i.parameters():
p.requires_grad = False
return idx
def get_groups(model, layer_groups):
group_indices = [len(g) for g in layer_groups]
curr_i = 0
group = []
for layer in model:
group_indices[curr_i] -= len(flatten_model(layer))
group.append(layer.__class__.__name__)
if group_indices[curr_i] == 0:
curr_i += 1
print(f'Group {curr_i}:', group)
group = []
def __init__(self, model, chip_size, num_classes):
super(AuxPSUnet, self).__init__()
self.model = model
for idx, i in enumerate(flatten_model(self.model)):
if hasattr(i, 'dilation'):
dilation = i.dilation
dilation = dilation[0] if isinstance(dilation, tuple) else dilation
if dilation > 1:
break
self.hook = hook_output(flatten_model(model)[idx - 1])
## returns the size of various activations
model_sizes(self.model, size=(chip_size, chip_size))
## Geting the stored parameters inside of the hook
aux_in_channels = self.hook.stored.shape[1]
del self.hook.stored
self.aux_logits = nn.Conv2d(aux_in_channels, num_classes, kernel_size=1)
def freeze(self):
"Freezes the pretrained backbone."
for idx, i in enumerate(flatten_model(self.learn.model)):
if hasattr(i, 'dilation'):
dilation = i.dilation
dilation = dilation[0] if isinstance(dilation, tuple) else dilation
if dilation > 1:
break
for p in i.parameters():
p.requires_grad = False
self.learn.layer_groups = split_model_idx(self.learn.model, [idx]) ## Could also call self.learn.freeze after this line because layer groups are now present.
def test_cs_learner_freeze(learn):
model = learn.model
for layer in fv.flatten_model(model.bu_body):
should_require = False
if isinstance(layer, fv.bn_types):
should_require = True
for p in layer.parameters():
assert p.requires_grad == should_require
branches = [model.td, model.laterals, model.emb]
if model.bu_head:
branches.append(model.bu_head)
for p in itertools.chain.from_iterable(b.parameters() for b in branches):
assert p.requires_grad
def _get_layers(model, cond=noop):
return [m for m in flatten_model(model) if cond(m)]
