File size: 1,053 Bytes
d4b83d9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
def replace_layer_recursive(model, old_layer, new_layer):
for name, layer in model._modules.items():
if layer == old_layer:
model._modules[name] = new_layer
return True
elif replace_layer_recursive(layer, old_layer, new_layer):
return True
return False
def replace_all_layer_type_recursive(model, old_layer_type, new_layer):
for name, layer in model._modules.items():
if isinstance(layer, old_layer_type):
model._modules[name] = new_layer
replace_all_layer_type_recursive(layer, old_layer_type, new_layer)
def find_layer_types_recursive(model, layer_types):
def predicate(layer):
return type(layer) in layer_types
return find_layer_predicate_recursive(model, predicate)
def find_layer_predicate_recursive(model, predicate):
result = []
for name, layer in model._modules.items():
if predicate(layer):
result.append(layer)
result.extend(find_layer_predicate_recursive(layer, predicate))
return result
|