wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-07-27 11:18:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

[fix] fix freeze top N layers

Browse Source
This commit is contained in:
theblackcat102
2022-12-31 17:43:27 +00:00
parent 918b7b7ec0
commit ba336fb087
3 changed files with 15 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/reward/instructor/utils.py
+4 -3
View File
@@ -41,23 +41,24 @@ def train_val_dataset(dataset, val_split=0.2):
return Subset(dataset, train_idx), Subset(dataset, val_idx)
def freeze_top_n_layers(model, target_layers):
for name, param in model.name_parameters():
for name, param in model.named_parameters():
if 'embed' in name:
param.requires_grad = False
elif 'layer' in name:
elif '.layer' in name:
tokens = name.split('.')
idx = 0
for token in tokens:
if 'layer' in token:
break
idx += 1
if idx >= len(tokens):
continue
layer_ = int(tokens[idx+1])
if layer_ < target_layers:
param.requires_grad = False
return model
def argument_parsing(parser):
default_params = {
'num_train_epochs': 4,