Yolov5s/ai_training/classification/load_lr_scheduler.py

11 lines
269 B
Python

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
def load_lr_scheduler(optimizer_ft, mode = 'max', patience=5):
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer_ft, mode = mode, patience=patience)
return scheduler