import numpy as np import tensorflow as tf tf.compat.v1.disable_eager_execution() tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR) import argparse, os, time, logging from tqdm import tqdm import pandas as pd import multiprocessing from functools import partial import pickle from model import UNet, ModelConfig from data_reader import DataReader_train, DataReader_test from postprocess import extract_picks, save_picks, save_picks_json, extract_amplitude, convert_true_picks, calc_performance from visulization import plot_waveform from util import EMA, LMA def read_args(): parser = argparse.ArgumentParser() parser.add_argument("--mode", default="train", help="train/train_valid/test/debug") parser.add_argument("--epochs", default=100, type=int, help="number of epochs (default: 10)") parser.add_argument("--batch_size", default=20, type=int, help="batch size") parser.add_argument("--learning_rate", default=0.01, type=float, help="learning rate") parser.add_argument("--drop_rate", default=0.0, type=float, help="dropout rate") parser.add_argument("--decay_step", default=-1, type=int, help="decay step") parser.add_argument("--decay_rate", default=0.9, type=float, help="decay rate") parser.add_argument("--momentum", default=0.9, type=float, help="momentum") parser.add_argument("--optimizer", default="adam", help="optimizer: adam, momentum") parser.add_argument("--summary", default=True, type=bool, help="summary") parser.add_argument("--class_weights", nargs="+", default=[1, 1, 1], type=float, help="class weights") parser.add_argument("--model_dir", default=None, help="Checkpoint directory (default: None)") parser.add_argument("--load_model", action="store_true", help="Load checkpoint") parser.add_argument("--log_dir", default="log", help="Log directory (default: log)") parser.add_argument("--num_plots", default=10, type=int, help="Plotting training results") parser.add_argument("--min_p_prob", default=0.3, type=float, help="Probability threshold for P pick") parser.add_argument("--min_s_prob", default=0.3, type=float, help="Probability threshold for S pick") parser.add_argument("--format", default="numpy", help="Input data format") parser.add_argument("--train_dir", default="./dataset/waveform_train/", help="Input file directory") parser.add_argument("--train_list", default="./dataset/waveform.csv", help="Input csv file") parser.add_argument("--valid_dir", default=None, help="Input file directory") parser.add_argument("--valid_list", default=None, help="Input csv file") parser.add_argument("--test_dir", default=None, help="Input file directory") parser.add_argument("--test_list", default=None, help="Input csv file") parser.add_argument("--result_dir", default="results", help="result directory") parser.add_argument("--plot_figure", action="store_true", help="If plot figure for test") parser.add_argument("--save_prob", action="store_true", help="If save result for test") args = parser.parse_args() return args def train_fn(args, data_reader, data_reader_valid=None): current_time = time.strftime("%y%m%d-%H%M%S") log_dir = os.path.join(args.log_dir, current_time) if not os.path.exists(log_dir): os.makedirs(log_dir) logging.info("Training log: {}".format(log_dir)) model_dir = os.path.join(log_dir, 'models') os.makedirs(model_dir) figure_dir = os.path.join(log_dir, 'figures') if not os.path.exists(figure_dir): os.makedirs(figure_dir) config = ModelConfig(X_shape=data_reader.X_shape, Y_shape=data_reader.Y_shape) if args.decay_step == -1: args.decay_step = data_reader.num_data // args.batch_size config.update_args(args) with open(os.path.join(log_dir, 'config.log'), 'w') as fp: fp.write('\n'.join("%s: %s" % item for item in vars(config).items())) with tf.compat.v1.name_scope('Input_Batch'): dataset = data_reader.dataset(args.batch_size, shuffle=True).repeat() batch = tf.compat.v1.data.make_one_shot_iterator(dataset).get_next() if data_reader_valid is not None: dataset_valid = data_reader_valid.dataset(args.batch_size, shuffle=False).repeat() valid_batch = tf.compat.v1.data.make_one_shot_iterator(dataset_valid).get_next() model = UNet(config, input_batch=batch) sess_config = tf.compat.v1.ConfigProto() sess_config.gpu_options.allow_growth = True # sess_config.log_device_placement = False with tf.compat.v1.Session(config=sess_config) as sess: summary_writer = tf.compat.v1.summary.FileWriter(log_dir, sess.graph) saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables(), max_to_keep=5) init = tf.compat.v1.global_variables_initializer() sess.run(init) if args.model_dir is not None: logging.info("restoring models...") latest_check_point = tf.train.latest_checkpoint(args.model_dir) saver.restore(sess, latest_check_point) if args.plot_figure: multiprocessing.set_start_method('spawn') pool = multiprocessing.Pool(multiprocessing.cpu_count()) flog = open(os.path.join(log_dir, 'loss.log'), 'w') train_loss = EMA(0.9) best_valid_loss = np.inf for epoch in range(args.epochs): progressbar = tqdm(range(0, data_reader.num_data, args.batch_size), desc="{}: epoch {}".format(log_dir.split("/")[-1], epoch)) for _ in progressbar: loss_batch, _, _ = sess.run([model.loss, model.train_op, model.global_step], feed_dict={model.drop_rate: args.drop_rate, model.is_training: True}) train_loss(loss_batch) progressbar.set_description("{}: epoch {}, loss={:.6f}, mean={:.6f}".format(log_dir.split("/")[-1], epoch, loss_batch, train_loss.value)) flog.write("epoch: {}, mean loss: {}\n".format(epoch, train_loss.value)) if data_reader_valid is not None: valid_loss = LMA() progressbar = tqdm(range(0, data_reader_valid.num_data, args.batch_size), desc="Valid:") for _ in progressbar: loss_batch, preds_batch, X_batch, Y_batch, fname_batch = sess.run([model.loss, model.preds, valid_batch[0], valid_batch[1], valid_batch[2]], feed_dict={model.drop_rate: 0, model.is_training: False}) valid_loss(loss_batch) progressbar.set_description("valid, loss={:.6f}, mean={:.6f}".format(loss_batch, valid_loss.value)) if valid_loss.value < best_valid_loss: best_valid_loss = valid_loss.value saver.save(sess, os.path.join(model_dir, "model_{}.ckpt".format(epoch))) flog.write("Valid: mean loss: {}\n".format(valid_loss.value)) else: loss_batch, preds_batch, X_batch, Y_batch, fname_batch = sess.run([model.loss, model.preds, batch[0], batch[1], batch[2]], feed_dict={model.drop_rate: 0, model.is_training: False}) saver.save(sess, os.path.join(model_dir, "model_{}.ckpt".format(epoch))) if args.plot_figure: pool.starmap( partial( plot_waveform, figure_dir=figure_dir, ), zip(X_batch, preds_batch, [x.decode() for x in fname_batch], Y_batch), ) # plot_waveform(X_batch, preds_batch, fname_batch, label=Y_batch, figure_dir=figure_dir) flog.flush() flog.close() return 0 def test_fn(args, data_reader): current_time = time.strftime("%y%m%d-%H%M%S") logging.info("{} log: {}".format(args.mode, current_time)) if args.model_dir is None: logging.error(f"model_dir = None!") return -1 if not os.path.exists(args.result_dir): os.makedirs(args.result_dir) figure_dir=os.path.join(args.result_dir, "figures") if not os.path.exists(figure_dir): os.makedirs(figure_dir) config = ModelConfig(X_shape=data_reader.X_shape, Y_shape=data_reader.Y_shape) config.update_args(args) with open(os.path.join(args.result_dir, 'config.log'), 'w') as fp: fp.write('\n'.join("%s: %s" % item for item in vars(config).items())) with tf.compat.v1.name_scope('Input_Batch'): dataset = data_reader.dataset(args.batch_size, shuffle=False) batch = tf.compat.v1.data.make_one_shot_iterator(dataset).get_next() model = UNet(config, input_batch=batch, mode='test') sess_config = tf.compat.v1.ConfigProto() sess_config.gpu_options.allow_growth = True # sess_config.log_device_placement = False with tf.compat.v1.Session(config=sess_config) as sess: saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables()) init = tf.compat.v1.global_variables_initializer() sess.run(init) logging.info("restoring models...") latest_check_point = tf.train.latest_checkpoint(args.model_dir) if latest_check_point is None: logging.error(f"No models found in model_dir: {args.model_dir}") return -1 saver.restore(sess, latest_check_point) flog = open(os.path.join(args.result_dir, 'loss.log'), 'w') test_loss = LMA() progressbar = tqdm(range(0, data_reader.num_data, args.batch_size), desc=args.mode) picks = [] true_picks = [] for _ in progressbar: loss_batch, preds_batch, X_batch, Y_batch, fname_batch, itp_batch, its_batch \ = sess.run([model.loss, model.preds, batch[0], batch[1], batch[2], batch[3], batch[4]], feed_dict={model.drop_rate: 0, model.is_training: False}) test_loss(loss_batch) progressbar.set_description("{}, loss={:.6f}, mean loss={:6f}".format(args.mode, loss_batch, test_loss.value)) picks_ = extract_picks(preds_batch, fname_batch) picks.extend(picks_) true_picks.extend(convert_true_picks(fname_batch, itp_batch, its_batch)) if args.plot_figure: plot_waveform(data_reader.config, X_batch, preds_batch, label=Y_batch, fname=fname_batch, itp=itp_batch, its=its_batch, figure_dir=figure_dir) save_picks(picks, args.result_dir) metrics = calc_performance(picks, true_picks, tol=3.0, dt=data_reader.config.dt) flog.write("mean loss: {}\n".format(test_loss)) flog.close() return 0 def main(args): logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO) coord = tf.train.Coordinator() if (args.mode == "train") or (args.mode == "train_valid"): with tf.compat.v1.name_scope('create_inputs'): data_reader = DataReader_train(format=args.format, data_dir=args.train_dir, data_list=args.train_list) if args.mode == "train_valid": data_reader_valid = DataReader_train(format=args.format, data_dir=args.valid_dir, data_list=args.valid_list) logging.info("Dataset size: train {}, valid {}".format(data_reader.num_data, data_reader_valid.num_data)) else: data_reader_valid = None logging.info("Dataset size: train {}".format(data_reader.num_data)) train_fn(args, data_reader, data_reader_valid) elif args.mode == "test": with tf.compat.v1.name_scope('create_inputs'): data_reader = DataReader_test(format=args.format, data_dir=args.test_dir, data_list=args.test_list) test_fn(args, data_reader) else: print("mode should be: train, train_valid, or test") return if __name__ == '__main__': args = read_args() main(args)