def __init__(self, config_fn, data_fn, nlat=96, nlon=144, nlev=30, ntime=48):
self.nlat, self.nlon = nlat, nlon
self.ngeo = nlat * nlon
repo_dir = os.getcwd().split('CBRAIN-CAM')[0] + 'CBRAIN-CAM/'
with open(config_fn, 'r') as f:
config = yaml.load(f)
self.model = tf.keras.models.load_model(
repo_dir + 'saved_models/' + config['exp_name'] + '/model.h5',
custom_objects={**layer_dict, **loss_dict})
out_scale_dict = load_pickle(config['output_dict'])
self.valid_gen = DataGenerator(
data_fn=config['data_dir'] + config['valid_fn'],
input_vars=config['inputs'],
output_vars=config['outputs'],
norm_fn=config['data_dir'] + config['norm_fn'],
input_transform=(config['input_sub'], config['input_div']),
output_transform=out_scale_dict,
batch_size=self.ngeo,
shuffle=False,
xarray=True,
var_cut_off=config['var_cut_off'] if 'var_cut_off' in config.keys() else None
)
def __init__(self, model, config_fn, data_fn, nlat=64, nlon=128, nlev=30, ntime=48):
self.nlat, self.nlon = nlat, nlon
self.ngeo = nlat * nlon
repo_dir = os.getcwd().split('CBRAIN-CAM')[0] + 'CBRAIN-CAM/'
with open(config_fn, 'r') as f:
config = yaml.load(f)
# self.model = tf.keras.models.load_model(
# repo_dir + 'saved_models/' + config['exp_name'] + '/model.h5',
# custom_objects={**layer_dict, **loss_dict})
# self.model = tf.keras.models.load_model(model_path,
# custom_objects={**layer_dict, **loss_dict})
self.model = model
out_scale_dict = load_pickle(config['output_dict'])
self.valid_gen = DataGenerator(
# tgb - 4/16/2019 - Changed data_fn to the argument data_fn of ModelDiagnostics
#data_fn=config['data_dir'] + config['valid_fn'],
data_fn = data_fn,
input_vars=config['inputs'],
output_vars=config['outputs'],
norm_fn=config['data_dir'] + config['norm_fn'],
input_transform=(config['input_sub'], config['input_div']),
output_transform=out_scale_dict,
batch_size=self.ngeo,
shuffle=False,
xarray=True,
var_cut_off=config['var_cut_off'] if 'var_cut_off' in config.keys() else None
)
def __init__(self,
model,
config_fn,
data_fn,
normalize_flag=True,
nlat=64,
nlon=128,
nlev=30,
ntime=48):
self.nlat, self.nlon = nlat, nlon
self.ngeo = nlat * nlon
repo_dir = os.getcwd().split('CBRAIN-CAM')[0] + 'CBRAIN-CAM/'
with open(config_fn, 'r') as f:
config = yaml.load(f)
self.model = model
out_scale_dict = load_pickle(config['output_dict'])
self.valid_gen = DataGenerator(
# tgb - 4/16/2019 - Changed data_fn to the argument data_fn of ModelDiagnostics
#data_fn=config['data_dir'] + config['valid_fn'],
data_fn=data_fn,
input_vars=config['inputs'],
output_vars=config['outputs'],
norm_fn=config['data_dir'] + config['norm_fn'],
input_transform=(config['input_sub'], config['input_div']),
output_transform=out_scale_dict,
batch_size=self.ngeo,
shuffle=False,
xarray=True,
normalize_flag=normalize_flag,
var_cut_off=config['var_cut_off']
if 'var_cut_off' in config.keys() else None)
def main(args):
"""Main training script."""
if args.gpu is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
limit_mem()
# Load output scaling dictionary
out_scale_dict = load_pickle(args.output_dict)
logging.info('Create training and validation data generators')
train_gen = DataGenerator(data_fn=args.data_dir + args.train_fn,
input_vars=args.inputs,
output_vars=args.outputs,
norm_fn=args.data_dir + args.norm_fn,
input_transform=(args.input_sub, args.input_div),
output_transform=out_scale_dict,
batch_size=args.batch_size,
shuffle=True,
var_cut_off=args.var_cut_off)
if args.valid_fn is not None:
valid_gen = DataGenerator(data_fn=args.data_dir + args.valid_fn,
input_vars=args.inputs,
output_vars=args.outputs,
norm_fn=args.data_dir + args.norm_fn,
input_transform=(args.input_sub,
args.input_div),
output_transform=out_scale_dict,
batch_size=args.batch_size * 10,
shuffle=False,
var_cut_off=args.var_cut_off)
else:
valid_gen = None
logging.info('Build model')
model = fc_model(input_shape=train_gen.n_inputs,
output_shape=train_gen.n_outputs,
hidden_layers=args.hidden_layers,
activation=args.activation,
conservation_layer=args.conservation_layer,
inp_sub=train_gen.input_transform.sub,
inp_div=train_gen.input_transform.div,
norm_q=out_scale_dict['PHQ'])
print(model.summary())
logging.info('Compile model')
if args.loss == 'weak_loss':
loss = WeakLoss(model.input,
inp_div=train_gen.input_transform.div,
inp_sub=train_gen.input_transform.sub,
norm_q=out_scale_dict['PHQ'],
alpha_mass=args.alpha_mass,
alpha_ent=args.alpha_ent,
noadiab=args.noadiab)
else:
loss = args.loss
metrics = [mse]
if args.conservation_metrics:
mass_loss = WeakLoss(model.input,
inp_div=train_gen.input_transform.div,
inp_sub=train_gen.input_transform.sub,
norm_q=out_scale_dict['PHQ'],
alpha_mass=1,
alpha_ent=0,
name='mass_loss',
noadiab=args.noadiab)
ent_loss = WeakLoss(model.input,
inp_div=train_gen.input_transform.div,
inp_sub=train_gen.input_transform.sub,
norm_q=out_scale_dict['PHQ'],
alpha_mass=0,
alpha_ent=1,
name='ent_loss',
noadiab=args.noadiab)
metrics += [mass_loss, ent_loss]
model.compile(args.optimizer, loss=loss, metrics=metrics)
lrs = LearningRateScheduler(LRUpdate(args.lr, args.lr_step,
args.lr_divide))
logging.info('Train model')
model.fit_generator(train_gen,
epochs=args.epochs,
validation_data=valid_gen,
callbacks=[lrs])
if args.exp_name is not None:
exp_dir = args.model_dir + args.exp_name + '/'
os.makedirs(exp_dir, exist_ok=True)
model_fn = exp_dir + 'model.h5'
logging.info(f'Saving model as {model_fn}')
model.save(model_fn)
if args.save_txt:
weights_fn = exp_dir + 'weights.h5'
logging.info(f'Saving weights as {weights_fn}')
model.save_weights(weights_fn)
save2txt(weights_fn, exp_dir)
save_norm(train_gen.input_transform, train_gen.output_transform,
exp_dir)
logging.info('Done!')
in_vars_RH = ['RH', 'TBP', 'PS', 'SOLIN', 'SHFLX', 'LHFLX']
#out_vars_RH = ['PHQ','TPHYSTND','FSNT', 'FSNS', 'FLNT', 'FLNS']
TRAINFILE_RH = 'CI_RH_M4K_NORM_train_shuffle.nc'
NORMFILE_RH = 'CI_RH_M4K_NORM_norm.nc'
VALIDFILE_RH = 'CI_RH_M4K_NORM_valid.nc' # Experiment 1/2
#VALIDFILE_RH = 'CI_RH_P4K_NORM_valid.nc' # Experiment 3/4
# In[8]:
train_gen_RH = DataGenerator(
data_fn=path + TRAINFILE_RH,
input_vars=in_vars_RH,
output_vars=out_vars_RH,
norm_fn=path + NORMFILE_RH,
input_transform=('mean', 'maxrs'),
output_transform=scale_dict_RH,
batch_size=BATCH_SIZE,
shuffle=True,
)
# ### Data Generator using TNS
# In[9]:
in_vars = ['QBP', 'TfromNS', 'PS', 'SOLIN', 'SHFLX', 'LHFLX']
#out_vars = ['PHQ','TPHYSTND','FSNT', 'FSNS', 'FLNT', 'FLNS']
TRAINFILE_TNS = 'CI_TNS_M4K_NORM_train_shuffle.nc'
NORMFILE_TNS = 'CI_TNS_M4K_NORM_norm.nc'
VALIDFILE_TNS = 'CI_TNS_M4K_NORM_valid.nc' # Experiment 1/2
scale_dict[
'RH'] = 0.01 * L_S / G, # Arbitrary 0.1 factor as specific humidity is generally below 2%
in_vars_RH = ['RH', 'TBP', 'PS', 'SOLIN', 'SHFLX', 'LHFLX']
out_vars_RH = ['PHQ', 'TPHYSTND', 'FSNT', 'FSNS', 'FLNT', 'FLNS']
TRAINFILE_RH = 'CI_RH_M4K_NORM_train_shuffle.nc'
NORMFILE_RH = 'CI_RH_M4K_NORM_norm.nc'
VALIDFILE_RH = 'CI_RH_M4K_NORM_valid.nc'
BASE_DIR = '/DFS-L/DATA/pritchard/ankitesg/'
train_gen_RH = DataGenerator(data_fn=f"{BASE_DIR}data/{TRAINFILE_RH}",
input_vars=in_vars_RH,
output_vars=out_vars_RH,
norm_fn=f"{BASE_DIR}data/{NORMFILE_RH}",
input_transform=('mean', 'maxrs'),
output_transform=scale_dict,
batch_size=1024,
shuffle=True,
normalize_flag=True)
in_vars = ['QBP', 'TfromNSV2', 'PS', 'SOLIN', 'SHFLX', 'LHFLX']
out_vars = ['PHQ', 'TPHYSTND', 'FSNT', 'FSNS', 'FLNT', 'FLNS']
TRAINFILE_TNS = 'CI_TNSV2_M4K_NORM_train_shuffle.nc'
NORMFILE_TNS = 'CI_TNSV2_M4K_NORM_norm.nc'
VALIDFILE_TNS = 'CI_TNSV2_M4K_NORM_valid.nc'
train_gen_TNS = DataGenerator(data_fn=f"{BASE_DIR}data/{TRAINFILE_TNS}",
input_vars=in_vars,
output_vars=out_vars,
def main(inargs):
"""Main function.
"""
# set GPU usage
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
keras.backend.tensorflow_backend.set_session(tf.Session(config=config))
if inargs.norm_fn is None:
norm_fn = inargs.train_fn.split('_shuffle')[0] + '_norm.nc'
else:
norm_fn = inargs.norm_fn
# Load train and valid set
train_gen = DataGenerator(
inargs.data_dir,
inargs.train_fn + '_features.nc',
inargs.train_fn + '_targets.nc',
inargs.batch_size,
norm_fn,
inargs.fsub, inargs.fdiv, inargs.tsub, inargs.tmult,
shuffle=True, noise=inargs.noise,
)
valid_gen = DataGenerator(
inargs.data_dir,
inargs.valid_fn + '_features.nc',
inargs.valid_fn + '_targets.nc',
16384, # Large batch size for speed!
norm_fn,
inargs.fsub, inargs.fdiv, inargs.tsub, inargs.tmult,
shuffle=False,
)
feature_shape = train_gen.feature_shape
target_shape = train_gen.target_shape
# Build and compile model
if inargs.convolution:
model = conv_model(
(30, 7) if inargs.tile else (30, 3),
4,
target_shape,
inargs.conv_layers,
inargs.hidden_layers,
inargs.lr,
loss_dict[inargs.loss],
kernel_size=inargs.kernel_size, stride=inargs.stride,
batch_norm=inargs.batch_norm,
activation=inargs.activation,
tile=inargs.tile,
locally_connected=inargs.locally_connected,
padding=inargs.padding,
dr=inargs.dr,
)
else: # Fully connected model
model = fc_model(
feature_shape,
target_shape,
inargs.hidden_layers,
inargs.lr,
loss_dict[inargs.loss],
batch_norm=inargs.batch_norm,
activation=inargs.activation,
dr=inargs.dr,
l2=inargs.l2,
partial_relu=inargs.partial_relu,
eq=inargs.eq,
fsub=train_gen.feature_norms[0],
fdiv=train_gen.feature_norms[1],
)
if inargs.verbose: print(model.summary())
callbacks_list = []
if inargs.log_dir is not None:
callbacks_list.append(TensorBoard(log_dir=inargs.log_dir +
inargs.exp_name))
if not inargs.lr_step == 0:
def lr_update(epoch):
# From goo.gl/GXQaK6
init_lr = inargs.lr
drop = 1./inargs.lr_divide
epochs_drop = inargs.lr_step
lr = init_lr * np.power(drop, np.floor((1+epoch)/epochs_drop))
print('lr:', lr)
return lr
callbacks_list.append(LearningRateScheduler(lr_update))
# Fit model
model.fit_generator(
train_gen.return_generator(inargs.convolution, inargs.tile),
train_gen.n_batches,
epochs=inargs.epochs - int(inargs.valid_after),
validation_data=None if inargs.valid_after
else valid_gen.return_generator(inargs.convolution, inargs.tile),
validation_steps=valid_gen.n_batches,
workers=inargs.n_workers,
max_queue_size=50,
callbacks=callbacks_list,
)
if inargs.valid_after:
# Run last epoch with validation
model.optimizer.lr = tf.Variable(lr_update(inargs.epochs))
if len(callbacks_list) == 1:
callbacks_list = []
else:
callbacks_list = [callbacks_list[0]] # No LR scheduler
model.fit_generator(
train_gen.return_generator(inargs.convolution, inargs.tile),
train_gen.n_batches,
epochs=1,
validation_data=valid_gen.return_generator(inargs.convolution, inargs.tile),
validation_steps=valid_gen.n_batches,
workers=inargs.n_workers,
max_queue_size=inargs.max_queue_size,
callbacks=callbacks_list,
)
if inargs.exp_name is not None:
if not os.path.exists(inargs.model_dir): os.mkdir(inargs.model_dir)
model.save(inargs.model_dir + '/' + inargs.exp_name + '.h5')