o
    1 i                     @   s  d dl Z d dlZd dlZd dlZd dlmZ d dlmZ G dd dejZ	e
dkre  Zejddd	d
 e \ZZejrDejdd dZededd ddgddZeje	ejddejr`dnddd ejeddddejedddddd ddd!Ze Zed"e j dS dS )#    N)tune)PopulationBasedTrainingc                   @   s8   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d ZdS )PBTBenchmarkExamplea  Toy PBT problem for benchmarking adaptive learning rate.

    The goal is to optimize this trainable's accuracy. The accuracy increases
    fastest at the optimal lr, which is a function of the current accuracy.

    The optimal lr schedule for this problem is the triangle wave as follows.
    Note that many lr schedules for real models also follow this shape:

     best lr
      ^
      |    /      |   /        |  /          | /            ------------> accuracy

    In this problem, using PBT with a population of 2-4 is sufficient to
    roughly approximate this lr schedule. Higher population sizes will yield
    faster convergence. Training will not converge without PBT.
    c                 C   s   |d | _ d| _d S )Nlrg        )r   accuracy)selfconfig r	   i/home/app/PaddleOCR-VL-test/.venv_paddleocr/lib/python3.10/site-packages/ray/tune/examples/pbt_example.pysetup#   s   

zPBTBenchmarkExample.setupc                 C   s   d}d}d}| j |k rd| j  | }ndd| j |  |  }tdtd|}t| j|t| j| }||k rD|  j d| t  7  _ n| j|krV|  j || t  8  _ |  j |tj  7  _ td| j | _ | j | j||| j |d kdS )	Nd         g{Gz?gMbP?g      ?r   )mean_accuracyZcur_lr
optimal_lrq_errdone)r   minmaxr   randomnpnormal)r   ZmidpointZq_toleranceZnoise_levelr   r   r	   r	   r
   step'   s(   

zPBTBenchmarkExample.stepc                 C   s   | j | jdS )Nr   r   r   )r   Zcheckpoint_dirr	   r	   r
   save_checkpointF   s   z#PBTBenchmarkExample.save_checkpointc                 C   s   |d | _ d S )Nr   )r   )r   
checkpointr	   r	   r
   load_checkpointL   s   z#PBTBenchmarkExample.load_checkpointc                 C   s   |d | _ || _dS )Nr   T)r   r   )r   Z
new_configr	   r	   r
   reset_configO   s   
z PBTBenchmarkExample.reset_configN)	__name__
__module____qualname____doc__r   r   r   r   r   r	   r	   r	   r
   r      s    r   __main__z--smoke-test
store_truezFinish quickly for testing)actionhelpr   )Znum_cpus   training_iterationc                   C   s   t ddS )N-C6?g{Gz?)r   uniformr	   r	   r	   r
   <lambda>e   s    r*      )r   Zsome_other_factor)Z	time_attrperturbation_intervalZhyperparam_mutationsZpbt_class_api_exampleT
   i  )r   r'   r      )Zcheckpoint_frequencyZcheckpoint_score_attributeZnum_to_keep)namestopverboseZcheckpoint_configr      )Z	schedulerZmetricmodeZreuse_actorsZnum_samplesr(   )Z
run_configZtune_configZparam_spacez!Best hyperparameters found were: )argparser   numpyr   Zrayr   Zray.tune.schedulersr   Z	Trainabler   r   ArgumentParserparseradd_argumentparse_known_argsargs_Z
smoke_testinitr,   ZpbtZTunerZ	RunConfigZCheckpointConfigZ
TuneConfigZtunerfitresultsprintZget_best_resultr   r	   r	   r	   r
   <module>   sf   H%