Spaces:
Runtime error
Runtime error
File size: 843 Bytes
9667e74 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
from torch import nn
def get_parameters(model, predicate):
for module in model.modules():
for param_name, param in module.named_parameters():
if predicate(module, param_name):
yield param
def get_parameters_conv(model, name):
return get_parameters(model, lambda m, p: isinstance(m, nn.Conv2d) and m.groups == 1 and p == name)
def get_parameters_conv_depthwise(model, name):
return get_parameters(model, lambda m, p: isinstance(m, nn.Conv2d)
and m.groups == m.in_channels
and m.in_channels == m.out_channels
and p == name)
def get_parameters_bn(model, name):
return get_parameters(model, lambda m, p: isinstance(m, nn.BatchNorm2d) and p == name)
|