Spaces:
Runtime error
Runtime error
File size: 1,354 Bytes
0c8d55e |
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 |
from transformers import Trainer
class UniVATrainer(Trainer):
def create_optimizer(self):
decay_parameters = self.get_decay_parameter_names(self.model)
optimizer_grouped_parameters = []
trainable_params = []
if self.optimizer is None:
for n, p in self.model.named_parameters():
if n in decay_parameters and p.requires_grad:
optimizer_grouped_parameters += [
{
"params": [p],
"weight_decay": self.args.weight_decay,
}
]
trainable_params.append(n)
elif n not in decay_parameters and p.requires_grad:
optimizer_grouped_parameters += [
{
"params": [p],
"weight_decay": 0.0,
}
]
trainable_params.append(n)
else:
p.requires_grad_(False)
optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(
self.args
)
self.optimizer = optimizer_cls(
optimizer_grouped_parameters, **optimizer_kwargs
)
return self.optimizer
|