Spaces:
Sleeping
Sleeping
File size: 1,712 Bytes
6cc8ae1 | 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 | """
Optimizer factory. Creates optimizers from config strings.
"""
import torch.optim as optim
import torch.nn as nn
def create_optimizer(
model_or_params,
optimizer_name: str = 'adam',
learning_rate: float = 0.001,
weight_decay: float = 0.0001,
momentum: float = 0.9,
) -> optim.Optimizer:
"""
Factory function to create optimizer.
Args:
model_or_params: nn.Module or list of parameters/param groups
optimizer_name: 'adam', 'adamw', 'sgd', or 'rmsprop'
learning_rate: base learning rate
weight_decay: L2 regularization
momentum: momentum for SGD/RMSProp
"""
# Get parameters
if isinstance(model_or_params, nn.Module):
params = model_or_params.parameters()
elif isinstance(model_or_params, list) and len(model_or_params) > 0:
if isinstance(model_or_params[0], dict):
# Already param groups
params = model_or_params
else:
params = model_or_params
else:
params = model_or_params
name = optimizer_name.lower()
if name == 'adam':
return optim.Adam(params, lr=learning_rate, weight_decay=weight_decay)
elif name == 'adamw':
return optim.AdamW(params, lr=learning_rate, weight_decay=weight_decay)
elif name == 'sgd':
return optim.SGD(params, lr=learning_rate, weight_decay=weight_decay,
momentum=momentum, nesterov=True)
elif name == 'rmsprop':
return optim.RMSprop(params, lr=learning_rate, weight_decay=weight_decay,
momentum=momentum)
else:
raise ValueError(f"Unknown optimizer: {name}. Use 'adam', 'adamw', 'sgd', or 'rmsprop'.")
|