File size: 2,723 Bytes
1ef0893 4346c95 1ef0893 4346c95 1ef0893 4346c95 1ef0893 4346c95 1ef0893 4346c95 1ef0893 4346c95 1ef0893 4346c95 1ef0893 4346c95 1ef0893 4346c95 |
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 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import torch
import os
# --- Configuration ---
# Pfad zur Eingabe-BIN-Datei
input_checkpoint_path = "./pytorch_model_oldLayerNames.bin"
# Pfad zur Ausgabe-BIN-Datei (wo die geänderte Version gespeichert wird)
output_checkpoint_path = "./pytorch_model_renamed_final.bin"
# Define the layers that *must* be named 'segmentation_head'
# These are the specific layers you identified
specific_segmentation_head_layers = [
'decode_head.conv_seg.bias',
'decode_head.conv_seg.weight',
'decode_head.convs.0.conv.bias',
'decode_head.convs.0.conv.weight',
'decode_head.convs.1.conv.bias',
'decode_head.convs.1.conv.weight',
'decode_head.convs.2.conv.bias',
'decode_head.convs.2.conv.weight',
'decode_head.convs.3.conv.bias',
'decode_head.convs.3.conv.weight',
'decode_head.fusion_conv.conv.bias',
'decode_head.fusion_conv.conv.weight'
]
# --- Check, ob die Eingabedatei existiert ---
if not os.path.exists(input_checkpoint_path):
print(f"Fehler: Eingabedatei nicht gefunden unter {input_checkpoint_path}. Bitte den Pfad korrigieren. ❌")
else:
# --- Checkpoint laden ---
state_dict = torch.load(input_checkpoint_path, map_location="cpu")
# --- Layer-Namen ändern und neues State Dict erstellen ---
new_state_dict = {}
renamed_count_segmentation = 0
renamed_count_segformer = 0
skipped_count = 0
for old_key, value in state_dict.items():
if old_key in specific_segmentation_head_layers:
# These specific layers get 'segmentation_head.' prefix
new_key = old_key.replace('decode_head.', 'segmentation_head.', 1)
new_state_dict[new_key] = value
renamed_count_segmentation += 1
elif old_key.startswith('decode_head.'):
# All other layers starting with 'decode_head.' get 'segformer_head.' prefix
new_key = old_key.replace('decode_head.', 'segformer_head.', 1)
new_state_dict[new_key] = value
renamed_count_segformer += 1
else:
# Keep other layers as they are (e.g., backbone layers)
new_state_dict[old_key] = value
skipped_count += 1
# --- Geändertes State Dict speichern ---
torch.save(new_state_dict, output_checkpoint_path)
print(f"✅ Fertig! Die umbenannte Datei wurde gespeichert unter: {output_checkpoint_path}")
print(f"Zusammenfassung der Umbenennungen:")
print(f" - '{renamed_count_segmentation}' Layer von 'decode_head.' zu 'segmentation_head.' umbenannt.")
print(f" - '{renamed_count_segformer}' Layer von 'decode_head.' zu 'segformer_head.' umbenannt.")
print(f" - '{skipped_count}' Layer behielten ihren ursprünglichen Namen (z.B. Backbone).") |