Spaces:
Running
Running
| from torch import nn | |
| import torch, copy | |
| import numpy as np | |
| def mlp( | |
| sizes, | |
| activation, | |
| output_activation=nn.Identity(), | |
| use_batchnorm=False, | |
| dropout_p=None, | |
| ): | |
| layers = [] | |
| for j in range(len(sizes) - 1): | |
| act = activation if j < len(sizes) - 2 else output_activation | |
| layers += [nn.Linear(sizes[j], sizes[j + 1])] | |
| if use_batchnorm: | |
| layers += nn.BatchNorm1d(sizes[j + 1]) | |
| layers += [act] | |
| if dropout_p is not None: | |
| layers += [nn.Dropout(p=dropout_p)] | |
| # if use_batchnorm: | |
| # layers += [nn.Linear(sizes[j], sizes[j+1]), nn.BatchNorm1d(sizes[j+1]), act] | |
| # else: layers += [nn.Linear(sizes[j], sizes[j+1]), act] | |
| return nn.Sequential(*layers) | |
| def get_scheduler(optimizer, scheduler_dict=None, epoch_num=None) -> torch.optim.lr_scheduler.LRScheduler: | |
| if scheduler_dict is None: | |
| raise ValueError("No scheduler_dict provided!") | |
| scheduler_type = scheduler_dict["type"] | |
| scheduler_dict.pop("type") | |
| if scheduler_type == "exponential": | |
| def lambda_rule(epoch): | |
| lr_l = max(scheduler_dict["min"], scheduler_dict["gamma"] ** epoch) | |
| return lr_l | |
| return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) | |
| # return torch.optim.lr_scheduler.ExponentialLR(optimizer, **scheduler_dict) | |
| elif scheduler_type == "linear": | |
| assert epoch_num != None, "linear schedule but epoch is None" | |
| def lambda_rule(epoch): | |
| lr_l = max( | |
| scheduler_dict["min"], 1 - epoch * scheduler_dict["slope"] | |
| ) # 1/float(epoch_num+1) | |
| return lr_l | |
| return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) | |
| elif scheduler_type == "Step_LR": | |
| print("Step_LR scheduler set") | |
| return torch.optim.lr_scheduler.StepLR(optimizer, **scheduler_dict) | |
| elif scheduler_type == "Step_LR_with_warmup": | |
| print("Step_LR_with_warmup scheduler set") | |
| warmup_epochs = scheduler_dict.pop("warmup_epochs", 5) | |
| # warm up phase | |
| warmup_scheduler = torch.optim.lr_scheduler.LinearLR( | |
| optimizer, start_factor=0.1, end_factor=1.0, total_iters=warmup_epochs | |
| ) | |
| # StepLR | |
| step_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, **scheduler_dict) | |
| # combine | |
| return torch.optim.lr_scheduler.SequentialLR( | |
| optimizer, schedulers=[warmup_scheduler, step_scheduler], milestones=[warmup_epochs] | |
| ) | |
| elif scheduler_type == "exponential_with_warmup": | |
| print("exponential_with_warmup scheduler set") | |
| warmup_epochs = scheduler_dict.pop("warmup_epochs", 5) | |
| # warm up phase | |
| warmup_scheduler = torch.optim.lr_scheduler.LinearLR( | |
| optimizer, start_factor=0.1, end_factor=1.0, total_iters=warmup_epochs | |
| ) | |
| # exponential decay | |
| def lambda_rule(epoch): | |
| lr_l = max(scheduler_dict["min"], scheduler_dict["gamma"] ** epoch) | |
| return lr_l | |
| exponential_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) | |
| # combine | |
| return torch.optim.lr_scheduler.SequentialLR( | |
| optimizer, schedulers=[warmup_scheduler, exponential_scheduler], milestones=[warmup_epochs] | |
| ) | |
| else: raise NotImplementedError(f"Scheduler {scheduler_type} is not implemented!") | |
| ################################################################################## | |
| # return torch.optim.lr_scheduler.StepLR(optimizer, 600, 0.5) | |
| # if scheduler_type == 'Plateau': | |
| # print('Plateau_LR shceduler set') | |
| # return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, **scheduler_dict) | |
| # return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5, verbose=True) | |
| def count_param(model): | |
| model_parameters = filter(lambda p: p.requires_grad, model.parameters()) | |
| num_params = sum([np.prod(p.size()) for p in model_parameters]) | |
| return num_params | |