def train(): model = UNet(cfg.input_shape) #编译和打印模型 model.compile(optimizer=cfg.optimizer, loss=cfg.loss, metrics=cfg.metrics) print_summary(model=model) #训练数据生成器G1 G1 = imageSegmentationGenerator(cfg.train_images, cfg.train_annotations, cfg.train_batch_size, cfg.n_classes, cfg.input_shape[0], cfg.input_shape[1], cfg.output_shape[0], cfg.output_shape[1]) #测试数据生成器G2 if cfg.validate: G2 = imageSegmentationGenerator(cfg.val_images, cfg.val_annotations, cfg.val_batch_size, cfg.n_classes, cfg.input_shape[0], cfg.input_shape[1], cfg.output_shape[0], cfg.output_shape[1]) #循环训练 save_index = 1 for ep in range(cfg.epochs): #1、训练两种方式 if not cfg.validate: #只有G1 hisroy = model.fit_generator( G1, steps_per_epoch=cfg.train_steps_per_epoch, workers=cfg.workers, epochs=1, verbose=1, use_multiprocessing=cfg.use_multiprocessing) else: #有G1和G2 hisroy = model.fit_generator( G1, steps_per_epoch=cfg.train_steps_per_epoch, workers=cfg.workers, epochs=1, verbose=1, use_multiprocessing=cfg.use_multiprocessing, validation_data=G2, validation_steps=cfg.validate_steps_per_epoch) # 2、保存模型 if save_index == cfg.epochs_save: save_index = 1 save_weights_name = 'model.{}'.format(ep) save_weights_path = os.path.join(cfg.save_weights_path, save_weights_name) model.save_weights(save_weights_path) save_index += 1
model = multi_gpu_model(model, gpus=args.gpus) # train_mode = False if args.mode == 'validation': print("inference") model.load_weights(filepath="outputs_4_channels/best_model_final.h5", by_name=True) model.evaluate() else: if args.weights is not None: model.load_weights(filepath=args.weigths, by_name=True, skip_mismatch=True) optimizer = optim.adam(lr=args.lr, decay=5e-4) model.compile(optimizer=optimizer, loss=bce_dice_loss, metrics=[dice_coeff]) log_filename = os.path.join(args.outputs, 'logger_model_train.csv') csv_log = CSVLogger(log_filename, separator=',', append=True) lr = ReduceLROnPlateau(monitor='val_loss', factor=0.8, patience=5, verbose=2, mode='min') ckpt_path = os.path.join( args.outputs, "best_weight_model_{epoch:03d}_{val_loss:.4f}.h5") ckpt = ModelCheckpoint(monitor='val_loss', filepath=ckpt_path,
def train_model(X_train, X_test, y_train, y_test, save_dir): """ Run train model process. Args: X_train (list): List of strings X_train files patches. X_test (list): List of strings X_test files patches. y_train (list): List of strings y_train files patches. y_test (list): List of strings y_test files patches. save_dir (pathlib.PosixPath): Path for save results. """ len_train, len_test = len(X_train), len(X_test) # Create datasets from data. train_dataset = tf.data.Dataset.from_tensor_slices( (X_train, y_train)).map(load_data, num_parallel_calls=AUTOTUNE) test_dataset = tf.data.Dataset.from_tensor_slices( (X_test, y_test)).map(load_data, num_parallel_calls=AUTOTUNE) # Prepare data for training. train_dataset = train_dataset.map(normalize, num_parallel_calls=AUTOTUNE) train_dataset = train_dataset.shuffle(len_train * 2) train_dataset = train_dataset.repeat() train_dataset = train_dataset.batch(BATCH_SIZE) train_dataset = train_dataset.prefetch(buffer_size=AUTOTUNE) test_dataset = test_dataset.map(normalize, num_parallel_calls=AUTOTUNE) test_dataset = test_dataset.batch(BATCH_SIZE) # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - # TRAIN MODEL. # Init paths. save_dir.mkdir(parents=True, exist_ok=True) path_weights = (save_dir / "weights.hdf5").as_posix() path_history = (save_dir / "training.csv").as_posix() path_model = (save_dir / "model").as_posix() path_predict_fulls = (save_dir / "predict_fulls.csv").as_posix() path_predict_means = (save_dir / "predict_means.csv").as_posix() # Init callbacks. callback_save_best_weights = tf.keras.callbacks.ModelCheckpoint( filepath=path_weights, monitor="val_loss", mode="min", save_best_only=True, verbose=0, ) callback_csv_logger = tf.keras.callbacks.CSVLogger( filename=path_history, separator=",", append=False, ) # Init optimizer. optimizer = tf.keras.optimizers.Adam(learning_rate=0.001) # Init model. model = UNet( input_shape=(IMAGE_NET_W, IMAGE_NET_H, CLASSES_COUNT), out_channels=CLASSES_COUNT, filters=[16, 32, 64, 128, 256], ) model.compile( optimizer=optimizer, loss=tf.keras.losses.MSE, metrics=[tf.keras.metrics.Accuracy()], ) # Fake predict for build model and summary. model.predict(next(iter(test_dataset))[0]) model.summary() # Load weights. if MODEL_LOAD_WEIGHTS: model.load_weights(path_weights) # Train model. if MODEL_RUN_TRAIN: model.fit( train_dataset, validation_data=test_dataset, epochs=EPOCHS, steps_per_epoch=len_train // BATCH_SIZE, validation_steps=len_test // BATCH_SIZE, callbacks=[callback_save_best_weights, callback_csv_logger], ) # Evaluate model. evaluate = model.evaluate(test_dataset) print(tabulate(zip(model.metrics_names, evaluate), tablefmt="fancy_grid")) # Make prediction. predict_dataset = model.predict(test_dataset) # Save model. model.save(path_model) # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - # COMPARE WITH REAL COORDINATES. X_test_names = [pathlib.Path(x).stem for x in X_test] result_array = np.zeros((len_test, CLASSES_COUNT)) for k, predict_array in enumerate(predict_dataset): for channel_type, channel_number in zip(CLASSES, range(CLASSES_COUNT)): # Get channel. predict_channel = predict_array[:, :, channel_number] # Get predicted coordinates with max value. maxes = np.argwhere(predict_channel.max() == predict_channel) x_p, y_p = np.mean(maxes, axis=0).astype(np.int) # Rescale predicted coordinates. x_p, y_p = rescale_coords( (x_p, y_p), src_size=(IMAGE_NET_W, IMAGE_NET_H), out_size=(IMAGE_SRC_W, IMAGE_SRC_H), ) # Load text coordinates. filename = PWD_COORDS / channel_type / f"{X_test_names[k]}.txt" x_t, y_t = np.fromfile(filename, sep=",") # Compute distance. # ! NOTE: in this dataset, 1 cm == 11 pix. distance = np.sqrt((x_t - x_p)**2 + (y_t - y_p)**2) / 11.0 # Save result. result_array[k][channel_number] = distance # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - # DUMP RESULTS. # Dump full results to CSV. results = pd.DataFrame(columns=["filename", *CLASSES]) for n, filename in enumerate(X_test_names): # Write rows of matrix to CSV. named_results = dict(zip(CLASSES, list(result_array[n]))) named_results["filename"] = filename results = results.append(named_results, ignore_index=True) results.to_csv(path_predict_fulls, index=False, sep=";") # Dump mean results to CSV. results = pd.DataFrame(columns=["class", "mean", "max"]) means = np.mean(result_array, axis=0) maxes = np.max(result_array, axis=0) for class_label, val_mean, val_max in zip(CLASSES, means, maxes): results = results.append( { "class": class_label, "mean": val_mean, "max": val_max }, ignore_index=True, ) results.to_csv(path_predict_means, index=False, sep=";") print(results)