Import Geant4 11.1.0.beta source tree
This commit is contained in:
Gabriele Cosmo
@@ -0,0 +1,134 @@
|
||||
"""
|
||||
** model **
|
||||
defines the VAE model class
|
||||
"""
|
||||
|
||||
# Setup
|
||||
import keras
|
||||
from tensorflow.keras.layers import Input, Dense, Lambda, Layer, Multiply, Add, concatenate
|
||||
from tensorflow.keras.layers import BatchNormalization
|
||||
from tensorflow.keras.models import Model
|
||||
from tensorflow.keras import backend as K
|
||||
from tensorflow.keras import metrics
|
||||
|
||||
# VAE model class
|
||||
class VAE:
|
||||
def __init__(self, **kwargs):
|
||||
self.original_dim = kwargs.get('original_dim')
|
||||
self.latent_dim = kwargs.get('latent_dim')
|
||||
self.batch_size = kwargs.get('batch_size')
|
||||
self.intermediate_dim1 = kwargs.get('intermediate_dim1')
|
||||
self.intermediate_dim2 = kwargs.get('intermediate_dim2')
|
||||
self.intermediate_dim3 = kwargs.get('intermediate_dim3')
|
||||
self.intermediate_dim4 = kwargs.get('intermediate_dim4')
|
||||
self.epsilon_std = kwargs.get('epsilon_std')
|
||||
self.mu = kwargs.get('mu')
|
||||
self.lr = kwargs.get('lr')
|
||||
self.epochs = kwargs.get('epochs')
|
||||
self.activ = kwargs.get('activ')
|
||||
self.outActiv = kwargs.get('outActiv')
|
||||
self.validation_split = kwargs.get('validation_split')
|
||||
self.wReco = kwargs.get('wReco')
|
||||
self.wkl = kwargs.get('wkl')
|
||||
self.optimizer = kwargs.get('optimizer')
|
||||
self.ki = kwargs.get('ki')
|
||||
self.bi = kwargs.get('bi')
|
||||
self.checkpoint_dir = kwargs.get('checkpoint_dir')
|
||||
self.earlyStop = kwargs.get('earlyStop')
|
||||
# KL divergence computation
|
||||
class KLDivergenceLayer(Layer):
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.is_placeholder = True
|
||||
super(KLDivergenceLayer, self).__init__(*args, **kwargs)
|
||||
def call(self, inputs):
|
||||
mu, log_var = inputs
|
||||
kl_batch = -self.wkl * K.sum(1 + log_var - K.square(mu) - K.exp(log_var), axis=-1)
|
||||
self.add_loss(K.mean(kl_batch), inputs=inputs)
|
||||
return inputs
|
||||
# Build the encoder
|
||||
xIn = Input((input_dim,))
|
||||
eCond = Input(shape=(1,))
|
||||
angleCond = Input(shape=(1,))
|
||||
GeoCond = Input(shape=(2,))
|
||||
mergedInput = concatenate([xIn, eCond, angleCond, GeoCond],)
|
||||
h1 = Dense(self.intermediate_dim1, activation=self.activ,
|
||||
kernel_initializer=self.ki, bias_initializer=self.bi)(mergedInput)
|
||||
h1 = BatchNormalization()(h1)
|
||||
h2 = Dense(self.intermediate_dim2, activation=self.activ,
|
||||
kernel_initializer=self.ki, bias_initializer=self.bi)(h1)
|
||||
h2 = BatchNormalization()(h2)
|
||||
h3 = Dense(self.intermediate_dim3, activation=self.activ,
|
||||
kernel_initializer=self.ki, bias_initializer=self.bi)(h2)
|
||||
h3 = BatchNormalization()(h3)
|
||||
h4 = Dense(self.intermediate_dim4, activation=self.activ,
|
||||
kernel_initializer=self.ki, bias_initializer=self.bi)(h3)
|
||||
h = BatchNormalization()(h4)
|
||||
z_mu = Dense(self.latent_dim,)(h)
|
||||
z_log_var = Dense(self.latent_dim,)(h)
|
||||
# compute the KL divergence
|
||||
z_mu, z_log_var = KLDivergenceLayer()([z_mu, z_log_var])
|
||||
# Reparameterization trick
|
||||
z_sigma = Lambda(lambda t: K.exp(.5*t))(z_log_var)
|
||||
eps = Input(tensor=K.random_normal(shape=(K.shape(xIn)[0], self.latent_dim)))
|
||||
z_eps = Multiply()([z_sigma, eps])
|
||||
z = Add()([z_mu, z_eps])
|
||||
zCond = concatenate([z,eCond,angleCond,GeoCond],)
|
||||
# This defines the encoder which takes noise and input and outputs the latent variable z
|
||||
self.encoder = Model(inputs=[xIn,eCond,angleCond,GeoCond,eps], outputs=zCond)
|
||||
# Build the decoder / Generator
|
||||
decoL4 = Dense(self.intermediate_dim4, input_dim=(self.latent_dim+4),
|
||||
activation=self.activ, kernel_initializer=self.ki, bias_initializer=self.bi)
|
||||
decoL4_BN = BatchNormalization()
|
||||
decoL3 = Dense(self.intermediate_dim3, input_dim=self.intermediate_dim4,
|
||||
activation=self.activ, kernel_initializer=self.ki, bias_initializer=self.bi)
|
||||
decoL3_BN = BatchNormalization()
|
||||
decoL2 = Dense(self.intermediate_dim2, input_dim=self.intermediate_dim3,
|
||||
activation=self.activ, kernel_initializer=self.ki, bias_initializer=self.bi)
|
||||
decoL2_BN = BatchNormalization()
|
||||
decoL1 = Dense(self.intermediate_dim1, input_dim=self.intermediate_dim2,
|
||||
activation=self.activ, kernel_initializer=self.ki, bias_initializer=self.bi)
|
||||
decoL1_BN = BatchNormalization()
|
||||
x_reco = Dense(self.original_dim, activation=self.outActiv)
|
||||
zDecoInput = Input(shape=(latent_dim+4,))
|
||||
x_recoDeco = x_reco((((decoL1_BN(decoL1(decoL2_BN(decoL2(decoL3_BN(decoL3(decoL4_BN(decoL4(zDecoInput))))))))))))
|
||||
# This defines the decoder which takes an input of size latent dimension + condition size dimension and outputs the reconstructed input version
|
||||
self.decoder = Model(inputs=[zDecoInput], outputs=[x_recoDeco])
|
||||
# This defines the reconstruction loss of the VAE model
|
||||
def reconstructionLoss(G4_Event, VAE_Event):
|
||||
return K.mean(self.wReco*K.sum(metrics.binary_crossentropy(G4_Event, VAE_Event)))
|
||||
# This defines the VAE model (encoder and decoder)
|
||||
self.vae = Model(inputs=[xIn,eCond,angleCond,GeoCond,eps], outputs=[self.decoder(self.encoder([xIn, eCond,angleCond,GeoCond,eps]))])
|
||||
self.vae.compile(optimizer=self.optimizer, loss=[reconstructionLoss] )
|
||||
# Training function
|
||||
def train(self, trainSet, eCond, angleCond, GeoCond):
|
||||
# If the early stopping flag is on then stop the training when a monitored metric (validation) has stopped improving after (patience) number of epochs
|
||||
if(self.earlyStop):
|
||||
from tensorflow.keras.callbacks import EarlyStopping
|
||||
cP = EarlyStopping(monitor='val_loss', min_delta=0.01, patience=5,verbose=1)
|
||||
# If the early stopping flag is off then run the training for the number of epochs and save the model every (period) epochs
|
||||
else:
|
||||
cP = keras.callbacks.ModelCheckpoint('%s/VAE-{epoch:02d}.h5'%self.checkpoint_dir, monitor='val_loss',
|
||||
verbose=0, save_best_only=False, save_weights_only=False, mode='auto',
|
||||
period=100)
|
||||
noise = np.random.normal(0,1, size = (trainSet.shape[0],latent_dim))
|
||||
history = self.vae.fit([trainSet, eCond, angleCond, GeoCond,noise], [trainSet],
|
||||
shuffle=True,
|
||||
epochs=self.epochs,
|
||||
verbose=1,
|
||||
validation_split=self.validation_split,
|
||||
batch_size=self.batch_size,
|
||||
callbacks=[cP]
|
||||
)
|
||||
return history
|
||||
# Encode function uses only the encoder to generate the latent representation of an input
|
||||
def encode(self, dataSet):
|
||||
return self.encoder.predict(dataSet, batch_size=self.batch_size)
|
||||
# Generate function uses only the decoder to generate new showers using the z_sample which is a vector of 10D Gaussians in addition to
|
||||
def generate(self, z_sample):
|
||||
return self.decoder.predict([z_sample])
|
||||
# Encode function
|
||||
def predict(self, dataSet):
|
||||
return self.vae.predict(dataSet, batch_size=self.batch_size)
|
||||
# Encode function
|
||||
def evaluate(self, dataSet):
|
||||
return self.vae.evaluate(dataSet, batch_size=self.batch_size)
|
||||
@@ -0,0 +1,87 @@
|
||||
"""
|
||||
** train **
|
||||
- defines data loading parameters and calls the data preprocessing function
|
||||
- defines the model parameters and instantiates the VAE model
|
||||
- performs the training
|
||||
"""
|
||||
|
||||
# 1. Data loading/preprocessing
|
||||
from utils import *
|
||||
# Directory where the HDF5 files are saved
|
||||
init_dir = './detector_'
|
||||
# Number of calorimeter layers
|
||||
nCells_z = 45
|
||||
# Segmentation in the r,phi direction
|
||||
nCells_r = 18
|
||||
nCells_phi = 50
|
||||
# Total number of readout cells (represents the number of nodes in the input/output layers of the model)
|
||||
original_dim = nCells_z*nCells_r*nCells_phi
|
||||
# Minimum and maximum primary particle energy to consider for training in GeV units
|
||||
min_energy = 1
|
||||
max_energy = 1024
|
||||
# Minimum and maximum primary particle angle to consider for training in degrees units
|
||||
min_angle = 50
|
||||
max_angle = 90
|
||||
# The preprocess function reads the data and performs preprocessing and encoding for the values of energy, angle and geometry
|
||||
energies_Train,condE_Train,condAngle_Train,condGeo_Train = preprocess(init_dir,original_dim,min_angle,max_angle,min_energy,max_energy)
|
||||
|
||||
# 2. Model architecture
|
||||
import model
|
||||
# Instantiate a VAE model and define all the parameters
|
||||
vae = model.VAE(batch_size=100 ,
|
||||
original_dim=original_dim,
|
||||
intermediate_dim1=100,
|
||||
intermediate_dim2=50,
|
||||
intermediate_dim3=20,
|
||||
intermediate_dim4=10+4,
|
||||
latent_dim=10,
|
||||
epsilon_std=1.,
|
||||
mu=0,
|
||||
epochs=10000,
|
||||
lr=0.001,
|
||||
activ=tf.keras.layers.LeakyReLU(),
|
||||
outActiv='sigmoid',
|
||||
validation_split=0.05,
|
||||
wReco=original_dim,
|
||||
wkl=0.5,
|
||||
optimizer=optimizers.Adam(),
|
||||
ki='RandomNormal',
|
||||
bi='Zeros',
|
||||
earlyStop=False,
|
||||
checkpoint_dir = "."
|
||||
)
|
||||
|
||||
# 3. Model training
|
||||
history = vae.train(energies_Train,
|
||||
condE_Train,
|
||||
condAngle_Train,
|
||||
condGeo_Train
|
||||
)
|
||||
|
||||
# 4. Save the model (ony the decoder part) after traing
|
||||
self.vae.decoder.save("decoder.h5")
|
||||
|
||||
# 5. Convert the model to ONNX format
|
||||
import keras2onnx
|
||||
import tensorflow
|
||||
# Create the Keras model and convert itinto an ONNX model
|
||||
kerasModel = tensorflow.keras.models.load_model("decoder.h5")
|
||||
onnxModel = keras2onnx.convert_keras(kerasModel,"name")
|
||||
# Save the ONNX model. Generator.onnx can then be used to perform the inference in the example
|
||||
keras2onnx.save_model(onnxModel,"Generator.onnx")
|
||||
|
||||
"""
|
||||
# In order to convert the model into a format that can be used with the LWTNN library
|
||||
# 1. After training :
|
||||
# serialize model to JSON
|
||||
json_model = self.vae.decoder.to_json()
|
||||
with open("decoder.json", "w") as json_file:
|
||||
json_file.write(json_model)
|
||||
# serialize weights to HDF5
|
||||
self.vae.decoder.save_weights("decoder.h5")
|
||||
# 2. Externally, after building the LWTNN code available at https://github.com/lwtnn/lwtnn
|
||||
# 2.1 Run the kerasfunc2json python script (available in lwtnn/ converters/) to generate a template file of your functional model input variables by calling:
|
||||
# $ kerasfunc2json.py decoder.json decoder.h5 > inputs.json
|
||||
# 2.2 Run again kerasfunc2json script to get your output file that would be used for the inference in the example
|
||||
# $ kerasfunc2json.py decoder.json decoder.h5 inputs.json > Generator.json
|
||||
"""
|
||||
@@ -0,0 +1,57 @@
|
||||
"""
|
||||
** utils **
|
||||
defines the data loading and preprocessing function
|
||||
"""
|
||||
|
||||
# Setup
|
||||
import h5py
|
||||
import numpy as np
|
||||
|
||||
# preprocess function returns the array of the shower energies and the condition arrays
|
||||
"""
|
||||
- init_dir: the name of the directory which contains the HDF5 files
|
||||
- size_1DVec: represents the size of the input and output layer of the VAE which corresponds to the total number of readout cells
|
||||
- min_energy,max_energy: minimum and maximum primary particle energy to consider for training in GeV units
|
||||
- min_angle and max_angle: minimum and maximum primary particle angle to consider for training in degrees units
|
||||
"""
|
||||
def preprocess(init_dir,size_1DVec,min_angle,max_angle,min_energy,max_energy):
|
||||
energies_Train = []
|
||||
condE_Train = []
|
||||
condAngle_Train = []
|
||||
condGeo_Train = []
|
||||
# This example is trained using 2 detector geometries
|
||||
for geo in [ 'SiW' , 'SciPb' ]:
|
||||
dirGeo = init_dir + geo + '/'
|
||||
energyParticle=min_energy
|
||||
# loop over the energies in powers of 2
|
||||
while(energyParticle<=max_energy):
|
||||
# loop over the angles in a step of 10
|
||||
for angleParticle in range(min_angle,max_angle+10,10):
|
||||
fName = 'Energy_%s_Angle_%s.hdf5' %(energyParticle,angleParticle)
|
||||
fName = dirGeo + fName
|
||||
# read the HDF5 file
|
||||
h5 = h5py.File(fName,'r')
|
||||
# get the key value of the group from the HDF5 file
|
||||
GroupKey = 'Grp_Angle_%s_E_%s'%(angleParticle,energyParticle)
|
||||
# get all key values of one group
|
||||
listKeys = list( h5[GroupKey].keys() )
|
||||
# loop over the events
|
||||
for ckey in listKeys:
|
||||
# scale the energy of each cell to the energy of the primary particle (in MeV units)
|
||||
energyArray = np.array(h5[GroupKey][ckey])/(energyParticle*1000)
|
||||
energies_Train.append( energyArray.reshape(size_1DVec) )
|
||||
# build the energy and angle condition vectors
|
||||
condE_Train.append( [energyParticle/mamax_energyxE]*len(listKeys) )
|
||||
condAngle_Train.append( [angleParticle/max_angle]*len(listKeys) )
|
||||
# build the geometry condition vector (1 hot encoding vector)
|
||||
if( geo == 'SiW' ):
|
||||
condGeo_Train.append( [[0,1]]*len(listKeys) )
|
||||
else:
|
||||
condGeo_Train.append( [[1,0]]*len(listKeys) )
|
||||
energyParticle*=2
|
||||
# return numpy arrays
|
||||
energies_Train = np.array(energies_Train)
|
||||
condE_Train = np.concatenate(condE_Train)
|
||||
condAngle_Train = np.concatenate(condAngle_Train)
|
||||
condGeo_Train = np.concatenate(condGeo_Train)
|
||||
return energies_Train,condE_Train,condAngle_Train,condGeo_Train
|
||||
Reference in New Issue
Block a user