Skip to content
Snippets Groups Projects
Select Git revision
  • Sprint/2021-10
  • master default protected
  • dev protected
  • Issue/3003-stsInstitute
  • gitkeep
  • Issue/2449-GuidPidSlugToProjectSettings
  • Issue/2309-docs
  • Fix/xxxx-updateDependencies
  • Issue/2364-testingKpiParser
  • Issue/2287-guestRole
  • Test/xxxx-pipelineTriggers
  • Issue/2102-gitLabResTypeRCV
  • Issue/2278-gitlabToS
  • Issue/2101-gitLabResTypeUi
  • Issue/1788-extractionCronjob
  • Issue/2183-kpiGeneratorResource
  • Issue/2222-resourceDateCreated
  • Issue/2221-projectDateCreated
  • Issue/1321-pidEnquiryOverhaul
  • Issue/1999-gitlabResourcesLib
  • Issue/1951-quotaImplementation
  • v2.22.0
  • v2.20.0
  • v2.19.1
  • v2.19.0
  • v2.18.0
  • v2.17.0
  • v2.16.2
  • v2.16.1
  • v2.16.0
  • v2.15.0
  • v2.14.0
  • v2.13.0
  • v2.12.1
  • v2.12.0
  • v2.11.1
  • v2.11.0
  • v2.10.1
  • v2.10.0
  • v2.9.1
  • v2.9.0
41 results

Program.cs

Blame
  • Code owners
    Assign users and groups as approvers for specific file changes. Learn more.
    train_MADE.py 4.01 KiB
    import torch
    from torch import tensor as T
    import numpy as np
    from sklearn.model_selection import train_test_split
    import h5py
    import argparse
    
    from distribution_estimation import models
    from distribution_estimation import training
    
    
    '''
    Authors: Nikolas Frediani (Nikolas.Frediani@physik.lmu.de), Michael Krämer, Philipp Mertsch, Kathrin Nippel
    Related publication: "SECRET: Stochasticity Emulator for Cosmic Ray Electrons", arXiv:2501.06011
    '''
    
    
    if __name__ == '__main__':
        '''
        This script trains MADE.
    
        It loads and prreprocesses the dataset, initializes the model and trainer,
        and trains the model for the specified number of epochs.
    
        Usage: python train_MADE.py -e <number_of_epochs> -i <input_dimension> -l <hidden_layers> -o <number_of_mixture_components>
    
        Arguments:
        -f, --filename: The filename of the training dataset (defaults to the provided file)
        -e, --n_epochs: Number of epochs to train for (required)
        -s, --start_epoch: Epoch to start at. If 0, a new model is initialized. If >0, try to load an existing model and continue training from that epoch (default: 0)
        -i, --input_dim: Input dimension (required)
        -l, --layer: Adds a hidden layer with the specified number of nodes (required, can be repeated for multiple layers)
        -o, --n_mixture: Number of mixture components in the model output. Used to calculate the output dimension (required)
        -m, --masks: Number of masks used (default: 1)
        -p, --model_path: Path where to look for saved models (default: 'models/')
    
        Example usage: python train_MADE.py -e 500 -i 41 -l 100 -l 100 -o 10 -m 1
        '''
        parser = argparse.ArgumentParser() # example default run parameters: ...python ...train_made.py -e 500 -i 41 -l 100 -o 10 -m 1
        parser.add_argument('-f','--filename',type=str,default='batches_regime_Klein-Nishina_SNrate_2.0e+04_dist_spiral_mode_full_causality_True_Ecut_1.0e+04_seeds_0-9999.hdf5',help='name of the file where data is stored')
        parser.add_argument('-e','--n_epochs',type=int,required=True,help='number of epochs to train for')
        parser.add_argument('-s','--start_epoch',type=int,default=0,help='epoch to start at. If 0, new model is initialised. If >0, try to load existing model and continue training from that epoch')
        parser.add_argument('-i','--input_dim',type=int,required=True,help='input dimension')
        parser.add_argument('-l','--layer',action='append',type=int,required=True,help='adds a hidden layers with _ number of nodes')
        parser.add_argument('-o','--n_mixture',type=int,required=True,help='number of mixture components in the model output. Used to calculate the output dimension')
        parser.add_argument('-m','--masks',type=int,default=1,help='number of masks used')
        parser.add_argument('-p','--model_path',type=str,default='runs/models/',choices=['runs/models/','runs/temp/'],help='path where to look for saved models') #choices restriction may be removed if necessary
        args = parser.parse_args()
    
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    
    
        # define energy bins
        Ens = torch.logspace(1., 5., 41)[14:33]
    
        # load data and standardise
        f = h5py.File('datasets/'+args.filename, 'r')
        print('dataset:',f['fluxes'])
        fluxes = (Ens**3 * T(np.array(f['fluxes']),dtype=torch.float)[14:33].t() * 0.06).log10()
        f.close()
    
        norm_fluxes = (fluxes-fluxes.mean(dim=0))/fluxes.std(dim=0).to(device)
    
        # prepare dataloaders
        train_set,test_set = train_test_split(norm_fluxes,train_size=0.9,random_state=151)
        train_data, val_data = torch.utils.data.DataLoader(train_set,50,shuffle=True), torch.utils.data.DataLoader(test_set,50,shuffle=True)
    
    
    
        # initialise model and trainer
        made = models.MADE(args.input_dim,args.layer,args.n_mixture,num_masks=args.masks).to(device)
        my_trainer = training.Trainer(made, training.logL_loss, torch.optim.Adam(made.parameters()), train_data, val_data)
        
        # load latest stage and train
        if args.start_epoch > 0:
            my_trainer.load_state(args.start_epoch,args.model_path)
        my_trainer.train(args.n_epochs)