Record output of an executing code periodically

24 Views Asked by At

I cannot find an answer which fits to my problem.

I have a function in python which calculate the energy of a physics system. To optimize the energy I apply SGD to this function. I can record the value of the energy at each SGD step, but I want to evaluate the performance of SGD. I don't need to calculate the total execution time of the code (I know how to do that), but I need to know the evolution of the energy during the execution of SGD. So for instance at each second I need to record the value of the energy calculated so far. I have a code like this

def run_one_SGD(max_steps, lr, precis, al_range, optim):
    al_0 = torch.distributions.Uniform(0, al_range).sample((size_basis,n_pairs)).double().requires_grad_(True)

    if optim == 'SGD':
        optimizer = torch.optim.SGD([al_0], lr = lr)
    elif optim == 'Adam':
        optimizer = torch.optim.Adam([al_0], lr = lr)

    en_hist = [1.0]
    time_hist = []
    i = 0
    precision = 1 
    start_time = time.perf_counter()

    while (i < max_steps and precision > precis):
        optimizer.zero_grad()
        S, K, P, H, e, eigs = find_en(al_0)
        e.backward()  
        optimizer.step()
        en_hist.append(e.clone().detach())
        i += 1
        print(f'At {optim} step {i} the energy is {e}')
        precision = abs((en_hist[-2] - en_hist[-1])/en_hist[-1])
        end_time = time.perf_counter()
        time_hist.insert(i, (end_time - start_time))
    return en_hist, time_hist

With this code I can record the execution time of each SGD step, but I would like to record the energy value e periodically (for instance each second) during the execution of the code.

Thanks in advance for help

Paolo

def run_one_SGD(max_steps, lr, precis, al_range, optim):
    al_0 = torch.distributions.Uniform(0, al_range).sample((size_basis,n_pairs)).double().requires_grad_(True)

    if optim == 'SGD':
        optimizer = torch.optim.SGD([al_0], lr = lr)
    elif optim == 'Adam':
        optimizer = torch.optim.Adam([al_0], lr = lr)

    en_hist = [1.0]
    time_hist = []
    i = 0
    precision = 1 
    start_time = time.perf_counter()

    while (i < max_steps and precision > precis):
        optimizer.zero_grad()
        S, K, P, H, e, eigs = find_en(al_0)
        e.backward()  
        optimizer.step()
        en_hist.append(e.clone().detach())
        i += 1
        print(f'At {optim} step {i} the energy is {e}')
        precision = abs((en_hist[-2] - en_hist[-1])/en_hist[-1])
        end_time = time.perf_counter()
        time_hist.insert(i, (end_time - start_time))
    return en_hist, time_hist
0

There are 0 best solutions below