def test_net(visualise, cache_scoremaps): logging.basicConfig(level=logging.INFO) cfg = load_config() dataset = create_dataset(cfg) dataset.set_shuffle(False) dataset.set_test_mode(True) sess, inputs, outputs = setup_pose_prediction(cfg) if cache_scoremaps: out_dir = cfg.scoremap_dir if not os.path.exists(out_dir): os.makedirs(out_dir) num_images = dataset.num_images predictions = np.zeros((num_images, ), dtype=np.object) for k in range(num_images): print("processing image {}/{}".format(k, num_images - 1)) batch = dataset.next_batch() outputs_np = sess.run(outputs, feed_dict={inputs: batch[Batch.inputs]}) scmap, locref = extract_cnn_output(outputs_np, cfg) pose = argmax_pose_predict(scmap, locref, cfg.stride) pose_refscale = np.copy(pose) pose_refscale[:, 0:2] /= cfg.global_scale predictions[k] = pose_refscale if visualise: img = np.squeeze(batch[Batch.inputs]).astype("uint8") visualize.show_heatmaps(cfg, img, scmap, pose) visualize.waitforbuttonpress() if cache_scoremaps: base = os.path.basename(batch[Batch.data_item].im_path) raw_name = os.path.splitext(base)[0] out_fn = os.path.join(out_dir, raw_name + ".mat") scipy.io.savemat(out_fn, mdict={"scoremaps": scmap.astype("float32")}) out_fn = os.path.join(out_dir, raw_name + "_locreg" + ".mat") if cfg.location_refinement: scipy.io.savemat( out_fn, mdict={"locreg_pred": locref.astype("float32")}) scipy.io.savemat("predictions.mat", mdict={"joints": predictions}) sess.close()
def train( config_yaml, displayiters, saveiters, maxiters, max_to_keep=5, keepdeconvweights=True, allow_growth=False, ): start_path = os.getcwd() os.chdir(str(Path(config_yaml).parents[0]) ) # switch to folder of config_yaml (for logging) setup_logging() cfg = load_config(config_yaml) net_type = cfg['net_type'] if cfg['dataset_type'] in ("scalecrop", "tensorpack", "deterministic"): print( "Switching batchsize to 1, as tensorpack/scalecrop/deterministic loaders do not support batches >1. Use imgaug/default loader." ) cfg["batch_size"] = 1 # in case this was edited for analysis.- dataset = create_dataset(cfg) batch_spec = get_batch_spec(cfg) batch, enqueue_op, placeholders = setup_preloading(batch_spec) losses = pose_net(cfg).train(batch) total_loss = losses["total_loss"] for k, t in losses.items(): TF.summary.scalar(k, t) merged_summaries = TF.summary.merge_all() if "snapshot" in Path(cfg['init_weights']).stem and keepdeconvweights: print("Loading already trained DLC with backbone:", net_type) variables_to_restore = slim.get_variables_to_restore() else: print("Loading ImageNet-pretrained", net_type) # loading backbone from ResNet, MobileNet etc. if "resnet" in net_type: variables_to_restore = slim.get_variables_to_restore( include=["resnet_v1"]) elif "mobilenet" in net_type: variables_to_restore = slim.get_variables_to_restore( include=["MobilenetV2"]) elif "efficientnet" in net_type: variables_to_restore = slim.get_variables_to_restore( include=["efficientnet"]) variables_to_restore = { var.op.name.replace("efficientnet/", "") + '/ExponentialMovingAverage': var for var in variables_to_restore } else: print("Wait for DLC 2.3.") restorer = TF.train.Saver(variables_to_restore) saver = TF.train.Saver( max_to_keep=max_to_keep ) # selects how many snapshots are stored, see https://github.com/AlexEMG/DeepLabCut/issues/8#issuecomment-387404835 if allow_growth == True: config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = TF.Session(config=config) else: sess = TF.Session() coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders) train_writer = TF.summary.FileWriter(cfg['log_dir'], sess.graph) if cfg.get("freezeencoder", False): if 'efficientnet' in net_type: print("Freezing ONLY supported MobileNet/ResNet currently!!") learning_rate, train_op, tstep = get_optimizer(total_loss, cfg) print("Freezing encoder...") learning_rate, _, train_op = get_optimizer_with_freeze(total_loss, cfg) else: learning_rate, train_op, tstep = get_optimizer(total_loss, cfg) sess.run(TF.global_variables_initializer()) sess.run(TF.local_variables_initializer()) # Restore variables from disk. restorer.restore(sess, cfg['init_weights']) if maxiters == None: max_iter = int(cfg['multi_step'][-1][1]) else: max_iter = min(int(cfg['multi_step'][-1][1]), int(maxiters)) # display_iters = max(1,int(displayiters)) print("Max_iters overwritten as", max_iter) if displayiters == None: display_iters = max(1, int(cfg['display_iters'])) else: display_iters = max(1, int(displayiters)) print("Display_iters overwritten as", display_iters) if saveiters == None: save_iters = max(1, int(cfg['save_iters'])) else: save_iters = max(1, int(saveiters)) print("Save_iters overwritten as", save_iters) cum_loss = 0.0 lr_gen = LearningRate(cfg) stats_path = Path(config_yaml).with_name("learning_stats.csv") lrf = open(str(stats_path), "w") print("Training parameter:") print(cfg) print("Starting training....") for it in range(max_iter + 1): if 'efficientnet' in net_type: dict = {tstep: it} current_lr = sess.run(learning_rate, feed_dict=dict) else: current_lr = lr_gen.get_lr(it) dict = {learning_rate: current_lr} [_, loss_val, summary] = sess.run( [train_op, total_loss, merged_summaries], feed_dict=dict, ) cum_loss += loss_val train_writer.add_summary(summary, it) if it % display_iters == 0 and it > 0: average_loss = cum_loss / display_iters cum_loss = 0.0 logging.info("iteration: {} loss: {} lr: {}".format( it, "{0:.4f}".format(average_loss), current_lr)) lrf.write("{}, {:.5f}, {}\n".format(it, average_loss, current_lr)) lrf.flush() # Save snapshot if (it % save_iters == 0 and it != 0) or it == max_iter: model_name = cfg['snapshot_prefix'] saver.save(sess, model_name, global_step=it) lrf.close() sess.close() coord.request_stop() coord.join([thread]) # return to original path. os.chdir(str(start_path))
def train(config_yaml, displayiters, saveiters, maxiters, max_to_keep=5, keepdeconvweights=True): start_path = os.getcwd() os.chdir(str(Path(config_yaml).parents[0]) ) #switch to folder of config_yaml (for logging) setup_logging() cfg = load_config(config_yaml) if cfg.dataset_type == 'default' or cfg.dataset_type == 'tensorpack' or cfg.dataset_type == 'deterministic': print( "Switching batchsize to 1, as default/tensorpack/deterministic loaders do not support batches >1. Use imgaug loader." ) cfg['batch_size'] = 1 #in case this was edited for analysis.- dataset = create_dataset(cfg) batch_spec = get_batch_spec(cfg) batch, enqueue_op, placeholders = setup_preloading(batch_spec) losses = pose_net(cfg).train(batch) total_loss = losses['total_loss'] for k, t in losses.items(): TF.summary.scalar(k, t) merged_summaries = TF.summary.merge_all() if 'snapshot' in Path(cfg.init_weights).stem and keepdeconvweights: print("Loading already trained DLC with backbone:", cfg.net_type) variables_to_restore = slim.get_variables_to_restore() else: print("Loading ImageNet-pretrained", cfg.net_type) #loading backbone from ResNet, MobileNet etc. if 'resnet' in cfg.net_type: variables_to_restore = slim.get_variables_to_restore( include=["resnet_v1"]) elif 'mobilenet' in cfg.net_type: variables_to_restore = slim.get_variables_to_restore( include=["MobilenetV2"]) else: print("Wait for DLC 2.3.") restorer = TF.train.Saver(variables_to_restore) saver = TF.train.Saver( max_to_keep=max_to_keep ) # selects how many snapshots are stored, see https://github.com/AlexEMG/DeepLabCut/issues/8#issuecomment-387404835 sess = TF.Session(config=config) coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders) train_writer = TF.summary.FileWriter(cfg.log_dir, sess.graph) learning_rate, train_op = get_optimizer(total_loss, cfg) sess.run(TF.global_variables_initializer()) sess.run(TF.local_variables_initializer()) # Restore variables from disk. restorer.restore(sess, cfg.init_weights) if maxiters == None: max_iter = int(cfg.multi_step[-1][1]) else: max_iter = min(int(cfg.multi_step[-1][1]), int(maxiters)) #display_iters = max(1,int(displayiters)) print("Max_iters overwritten as", max_iter) if displayiters == None: display_iters = max(1, int(cfg.display_iters)) else: display_iters = max(1, int(displayiters)) print("Display_iters overwritten as", display_iters) if saveiters == None: save_iters = max(1, int(cfg.save_iters)) else: save_iters = max(1, int(saveiters)) print("Save_iters overwritten as", save_iters) cum_loss = 0.0 lr_gen = LearningRate(cfg) stats_path = Path(config_yaml).with_name('learning_stats.csv') lrf = open(str(stats_path), 'w') print("Training parameter:") print(cfg) print("Starting training....") for it in range(max_iter + 1): current_lr = lr_gen.get_lr(it) [_, loss_val, summary] = sess.run([train_op, total_loss, merged_summaries], feed_dict={learning_rate: current_lr}) cum_loss += loss_val train_writer.add_summary(summary, it) if it % display_iters == 0 and it > 0: average_loss = cum_loss / display_iters cum_loss = 0.0 logging.info("iteration: {} loss: {} lr: {}".format( it, "{0:.4f}".format(average_loss), current_lr)) lrf.write("{}, {:.5f}, {}\n".format(it, average_loss, current_lr)) lrf.flush() # Save snapshot if (it % save_iters == 0 and it != 0) or it == max_iter: model_name = cfg.snapshot_prefix saver.save(sess, model_name, global_step=it) lrf.close() sess.close() coord.request_stop() coord.join([thread]) #return to original path. os.chdir(str(start_path))
def train( config_yaml, displayiters, saveiters, maxiters, max_to_keep=5, keepdeconvweights=True, allow_growth=False, ): start_path = os.getcwd() os.chdir( str(Path(config_yaml).parents[0]) ) # switch to folder of config_yaml (for logging) setup_logging() cfg = load_config(config_yaml) if cfg["optimizer"] != "adam": print( "Setting batchsize to 1! Larger batchsize not supported for this loader:", cfg["dataset_type"], ) cfg["batch_size"] = 1 if ( cfg["partaffinityfield_predict"] and "multi-animal" in cfg["dataset_type"] ): # the PAF code currently just hijacks the pairwise net stuff (for the batch feeding via Batch.pairwise_targets: 5) print("Activating limb prediction...") cfg["pairwise_predict"] = True dataset = create_dataset(cfg) batch_spec = get_batch_spec(cfg) batch, enqueue_op, placeholders = setup_preloading(batch_spec) losses = pose_net(cfg).train(batch) total_loss = losses["total_loss"] for k, t in losses.items(): TF.summary.scalar(k, t) merged_summaries = TF.summary.merge_all() net_type = cfg["net_type"] if "snapshot" in Path(cfg["init_weights"]).stem and keepdeconvweights: print("Loading already trained DLC with backbone:", net_type) variables_to_restore = slim.get_variables_to_restore() else: print("Loading ImageNet-pretrained", net_type) # loading backbone from ResNet, MobileNet etc. if "resnet" in net_type: variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"]) elif "mobilenet" in net_type: variables_to_restore = slim.get_variables_to_restore( include=["MobilenetV2"] ) elif "efficientnet" in net_type: variables_to_restore = slim.get_variables_to_restore( include=["efficientnet"] ) variables_to_restore = { var.op.name.replace("efficientnet/", "") + "/ExponentialMovingAverage": var for var in variables_to_restore } else: print("Wait for DLC 2.3.") restorer = TF.train.Saver(variables_to_restore) saver = TF.train.Saver( max_to_keep=max_to_keep ) # selects how many snapshots are stored, see https://github.com/AlexEMG/DeepLabCut/issues/8#issuecomment-387404835 if allow_growth: config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = TF.Session(config=config) else: sess = TF.Session() coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders) train_writer = TF.summary.FileWriter(cfg["log_dir"], sess.graph) learning_rate, train_op, tstep = get_optimizer(total_loss, cfg) sess.run(TF.global_variables_initializer()) sess.run(TF.local_variables_initializer()) restorer.restore(sess, cfg["init_weights"]) if maxiters == None: max_iter = int(cfg["multi_step"][-1][1]) else: max_iter = min(int(cfg["multi_step"][-1][1]), int(maxiters)) # display_iters = max(1,int(displayiters)) print("Max_iters overwritten as", max_iter) if displayiters == None: display_iters = max(1, int(cfg["display_iters"])) else: display_iters = max(1, int(displayiters)) print("Display_iters overwritten as", display_iters) if saveiters == None: save_iters = max(1, int(cfg["save_iters"])) else: save_iters = max(1, int(saveiters)) print("Save_iters overwritten as", save_iters) cumloss, partloss, locrefloss, pwloss = 0.0, 0.0, 0.0, 0.0 lr_gen = LearningRate(cfg) stats_path = Path(config_yaml).with_name("learning_stats.csv") lrf = open(str(stats_path), "w") print("Training parameters:") print(cfg) print("Starting multi-animal training....") for it in range(max_iter + 1): if "efficientnet" in net_type: dict = {tstep: it} current_lr = sess.run(learning_rate, feed_dict=dict) else: current_lr = lr_gen.get_lr(it) dict = {learning_rate: current_lr} # [_, loss_val, summary] = sess.run([train_op, total_loss, merged_summaries],feed_dict={learning_rate: current_lr}) [_, alllosses, loss_val, summary] = sess.run( [train_op, losses, total_loss, merged_summaries], feed_dict=dict ) partloss += alllosses["part_loss"] # scoremap loss if cfg["location_refinement"]: locrefloss += alllosses["locref_loss"] if cfg["pairwise_predict"]: # paf loss pwloss += alllosses["pairwise_loss"] cumloss += loss_val train_writer.add_summary(summary, it) if it % display_iters == 0 and it > 0: logging.info( "iteration: {} loss: {} scmap loss: {} locref loss: {} limb loss: {} lr: {}".format( it, "{0:.4f}".format(cumloss / display_iters), "{0:.4f}".format(partloss / display_iters), "{0:.4f}".format(locrefloss / display_iters), "{0:.4f}".format(pwloss / display_iters), current_lr, ) ) lrf.write( "iteration: {}, loss: {}, scmap loss: {}, locref loss: {}, limb loss: {}, lr: {}\n".format( it, "{0:.4f}".format(cumloss / display_iters), "{0:.4f}".format(partloss / display_iters), "{0:.4f}".format(locrefloss / display_iters), "{0:.4f}".format(pwloss / display_iters), current_lr, ) ) cumloss, partloss, locrefloss, pwloss = 0.0, 0.0, 0.0, 0.0 lrf.flush() # Save snapshot if (it % save_iters == 0 and it != 0) or it == max_iter: model_name = cfg["snapshot_prefix"] saver.save(sess, model_name, global_step=it) lrf.close() sess.close() coord.request_stop() coord.join([thread]) # return to original path. os.chdir(str(start_path))
def train(config_yaml, displayiters, saveiters, maxiters, max_to_keep=5): start_path = os.getcwd() os.chdir(str(Path(config_yaml).parents[0]) ) #switch to folder of config_yaml (for logging) setup_logging() cfg = load_config(config_yaml) cfg['batch_size'] = 1 #in case this was edited for analysis. dataset = create_dataset(cfg) batch_spec = get_batch_spec(cfg) batch, enqueue_op, placeholders = setup_preloading(batch_spec) losses = pose_net(cfg).train(batch) total_loss = losses['total_loss'] for k, t in losses.items(): TF.summary.scalar(k, t) merged_summaries = TF.summary.merge_all() variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"]) restorer = TF.train.Saver(variables_to_restore) saver = TF.train.Saver( max_to_keep=max_to_keep ) # selects how many snapshots are stored, see https://github.com/AlexEMG/DeepLabCut/issues/8#issuecomment-387404835 # sess = TF.Session() sess = TF.Session(config=TF.ConfigProto(device_count={'GPU': 0})) coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders) train_writer = TF.summary.FileWriter(cfg.log_dir, sess.graph) learning_rate, train_op = get_optimizer(total_loss, cfg) sess.run(TF.global_variables_initializer()) sess.run(TF.local_variables_initializer()) # Restore variables from disk. restorer.restore(sess, cfg.init_weights) if maxiters == None: max_iter = int(cfg.multi_step[-1][1]) else: max_iter = min(int(cfg.multi_step[-1][1]), int(maxiters)) #display_iters = max(1,int(displayiters)) print("\n\nMax_iters overwritten as", max_iter) if displayiters == None: display_iters = max(1, int(cfg.display_iters)) else: display_iters = max(1, int(displayiters)) print("Display_iters overwritten as", display_iters) if saveiters == None: save_iters = max(1, int(cfg.save_iters)) else: save_iters = max(1, int(saveiters)) print("Save_iters overwritten as", save_iters) cum_loss = 0.0 lr_gen = LearningRate(cfg) stats_path = Path(config_yaml).with_name('learning_stats.csv') lrf = open(str(stats_path), 'w') print("\nTraining parameter:\n") pprint.pprint(cfg) print("\n\nStarting training....") start = time.time() print("\nStarting time of training: {} \n".format( datetime.datetime.now())) for it in range(max_iter + 1): current_lr = lr_gen.get_lr(it) [_, loss_val, summary] = sess.run([train_op, total_loss, merged_summaries], feed_dict={learning_rate: current_lr}) cum_loss += loss_val train_writer.add_summary(summary, it) if it % display_iters == 0: end = time.time() hours, rem = divmod(end - start, 3600) time_hours, time_rem = divmod(end, 3600) minutes, seconds = divmod(rem, 60) time_mins, _ = divmod(time_rem, 60) average_loss = cum_loss / display_iters cum_loss = 0.0 logging.info( "iteration: {}/{}, loss: {:.4f}, lr: {}, | Elapsed Time: {:0>2}:{:0>2}:{:05.2f}, Time: {}" .format(it, max_iter, average_loss, current_lr, int(hours), int(minutes), seconds, datetime.datetime.now().strftime("%H:%M"))) lrf.write("{}, {:.5f}, {}\n".format(it, average_loss, current_lr)) lrf.flush() # Save snapshot if (it % save_iters == 0 and it != 0) or it == max_iter: model_name = cfg.snapshot_prefix saver.save(sess, model_name, global_step=it) lrf.close() sess.close() coord.request_stop() coord.join([thread]) #return to original path. os.chdir(str(start_path))
def train(config_yaml, displayiters, saveiters, max_to_keep=5): start_path = os.getcwd() os.chdir(str(Path(config_yaml).parents[0]) ) #switch to folder of config_yaml (for logging) setup_logging() cfg = load_config(config_yaml) cfg['batch_size'] = 1 #in case this was edited for analysis. dataset = create_dataset(cfg) batch_spec = get_batch_spec(cfg) batch, enqueue_op, placeholders = setup_preloading(batch_spec) losses = pose_net(cfg).train(batch) total_loss = losses['total_loss'] for k, t in losses.items(): tf.summary.scalar(k, t) merged_summaries = tf.summary.merge_all() variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"]) restorer = tf.train.Saver(variables_to_restore) saver = tf.train.Saver( max_to_keep=max_to_keep ) # selects how many snapshots are stored, see https://github.com/AlexEMG/DeepLabCut/issues/8#issuecomment-387404835 sess = tf.Session() coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders) train_writer = tf.summary.FileWriter(cfg.log_dir, sess.graph) learning_rate, train_op = get_optimizer(total_loss, cfg) sess.run(tf.global_variables_initializer()) sess.run(tf.local_variables_initializer()) # Restore variables from disk. restorer.restore(sess, cfg.init_weights) max_iter = int(cfg.multi_step[-1][1]) if displayiters == None: display_iters = max(1, int(cfg.display_iters)) else: display_iters = max(1, int(displayiters)) print("Display_iters overwritten as", display_iters) if saveiters == None: save_iters = max(1, int(cfg.save_iters)) else: save_iters = max(1, int(saveiters)) print("Save_iters overwritten as", save_iters) cum_loss = 0.0 lr_gen = LearningRate(cfg) stats_path = Path(config_yaml).with_name('learning_stats.csv') lrf = open(str(stats_path), 'w') print("Training parameter:") print(cfg) print("Starting training....") for it in range(max_iter + 1): current_lr = lr_gen.get_lr(it) [_, loss_val, summary] = sess.run([train_op, total_loss, merged_summaries], feed_dict={learning_rate: current_lr}) cum_loss += loss_val train_writer.add_summary(summary, it) if it % display_iters == 0 and it > 0: average_loss = cum_loss / display_iters cum_loss = 0.0 logging.info("iteration: {} loss: {} lr: {}".format( it, "{0:.4f}".format(average_loss), current_lr)) lrf.write("{}, {:.5f}, {}\n".format(it, average_loss, current_lr)) lrf.flush() # Save snapshot if (it % save_iters == 0 and it != 0) or it == max_iter: model_name = cfg.snapshot_prefix saver.save(sess, model_name, global_step=it) lrf.close() sess.close() coord.request_stop() coord.join([thread]) #return to original path. os.chdir(str(start_path))
def train(cfg, pose_config_yaml, displayiters, saveiters, maxiters, max_to_keep=5): start_path = os.getcwd() os.chdir(str(Path(pose_config_yaml).parents[0]) ) #switch to folder of config_yaml (for logging) setup_logging() pose_cfg = load_config(pose_config_yaml) pose_cfg['batch_size'] = 1 #in case this was edited for analysis. # TODO:: Cleanup (Setting up validation) early_stopping_thresh = 50 validator = Validator(cfg, pose_cfg, pose_config_yaml) dataset = create_dataset(pose_cfg) batch_spec = get_batch_spec(pose_cfg) batch, enqueue_op, placeholders = setup_preloading(batch_spec) losses = (pose_net.PoseNet(pose_cfg)).train(batch) total_loss = losses['total_loss'] for k, t in losses.items(): tf.summary.scalar(k, t) merged_summaries = tf.summary.merge_all() variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"]) restorer = tf.train.Saver(variables_to_restore) saver = tf.train.Saver( max_to_keep=max_to_keep ) # selects how many snapshots are stored, see https://github.com/AlexEMG/DeepLabCut/issues/8#issuecomment-387404835 sess = tf.Session() coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders) train_writer = tf.summary.FileWriter(pose_cfg.log_dir, sess.graph) learning_rate, train_op = get_optimizer(total_loss, pose_cfg) sess.run(tf.global_variables_initializer()) sess.run(tf.local_variables_initializer()) # Restore variables from disk. if pose_cfg.init_weights == 'He': # Default in ResNet print("Random weight initalization using He.") else: print("Pretrained weight initalization.") restorer.restore(sess, pose_cfg.init_weights) if maxiters == None: max_iter = int(pose_cfg.multi_step[-1][1]) else: max_iter = min(int(pose_cfg.multi_step[-1][1]), int(maxiters)) #display_iters = max(1,int(displayiters)) print("Max_iters overwritten as", max_iter) if displayiters == None: display_iters = max(1, int(pose_cfg.display_iters)) else: display_iters = max(1, int(displayiters)) print("Display_iters overwritten as", display_iters) if saveiters == None: save_iters = max(1, int(pose_cfg.save_iters)) else: save_iters = max(1, int(saveiters)) print("Save_iters overwritten as", save_iters) # Visualize first layer # import numpy as np # from skimage import color # vars = tf.trainable_variables() # print(vars) # vars_vals = sess.run(vars[0]) # vars_vals = np.moveaxis(vars_vals, -1, 0) # #vars_vals = (vars_vals - np.amin(vars_vals)) / (np.amax(vars_vals) - np.amin(vars_vals)) # vars_vals = color.rgb2gray(vars_vals) # disp_heatmap(vars_vals, cols=8, title='context2_pe300') cum_loss = 0.0 lr_gen = LearningRate(pose_cfg) validerror_min = float('inf') last_min = 0 stats_path = Path(pose_config_yaml).with_name('learning_stats.csv') lrf = open(str(stats_path), 'w') print("Training parameter:") print(pose_cfg) print("Starting training....") for it in range(max_iter + 1): current_lr = lr_gen.get_lr(it) [_, loss_val, summary] = sess.run([train_op, total_loss, merged_summaries], feed_dict={learning_rate: current_lr}) cum_loss += loss_val train_writer.add_summary(summary, it) if it % display_iters == 0 and it > 0: average_loss = cum_loss / display_iters cum_loss = 0.0 logging.info("iteration: {} loss: {} lr: {}".format( it, "{0:.4f}".format(average_loss), current_lr)) lrf.write("{}, {:.5f}, {}\n".format(it, average_loss, current_lr)) lrf.flush() # Save snapshot if (it % save_iters == 0 and it != 0) or it == max_iter: model_name = pose_cfg.snapshot_prefix print("Calculating validation performance...") validerror = validator.validate(sess, it) if validerror < validerror_min: validerror_min = validerror last_min = 0 saver.save(sess, model_name, global_step=it) else: last_min += 1 if last_min > early_stopping_thresh: print( "Early stopping because early_stopping_thresh has been exceeded." ) break lrf.close() validator.close() sess.close() coord.request_stop() coord.join([thread]) #return to original path. os.chdir(str(start_path))
def store_prediction_layer(task, date, shuffle, overwrite_snapshot=None, allow_growth=True): #%% from deeplabcut.pose_estimation_tensorflow.dataset.factory import ( create as create_dataset, ) from deeplabcut.pose_estimation_tensorflow.dataset.pose_dataset import Batch from deeplabcut.pose_estimation_tensorflow.nnet.net_factory import pose_net # from deeplabcut.pose_estimation_tensorflow.train import ( get_batch_spec, setup_preloading, start_preloading, get_optimizer, LearningRate, ) from deepgraphpose.PoseDataLoader import DataLoader from deepgraphpose.utils_model import load_dlc_snapshot, get_train_config, \ get_model_config #%% data_info = DataLoader(task) #%% cfg = get_model_config(task, data_info.model_data_dir, scorer=data_info.scorer, date=date) #%% dlc_cfg = get_train_config(cfg, shuffle) trainingsnapshot_name, trainingsnapshot, dlc_cfg = load_dlc_snapshot( dlc_cfg, overwrite_snapshot=overwrite_snapshot) #%% # Batch is a class filled with indices TF.reset_default_graph() # create dataset dataset = create_dataset(dlc_cfg) #%% # train: inputs, part_score_targets, part_score_weights, locref_mask batch_spec = get_batch_spec(dlc_cfg) # queing batch, enqueue_op, placeholders = setup_preloading(batch_spec) # init graph pn = pose_net(dlc_cfg) # extracts features, and runs it through a covnet, inputs = batch[Batch.inputs] net, end_points = pn.extract_features(inputs) # net is the input to the conv2d_transpose layer heads = pn.prediction_layers(net, end_points) #%% multi_class_labels = batch[Batch.part_score_targets] weigh_part_predictions = dlc_cfg.weigh_part_predictions part_score_weights = batch[ Batch.part_score_weights] if weigh_part_predictions else 1.0 #%% from deeplabcut.pose_estimation_tensorflow.nnet import losses #%% def add_part_loss(multi_class_labels, logits, part_score_weights): return tf.losses.sigmoid_cross_entropy(multi_class_labels, logits, part_score_weights) loss = {} logits = heads['part_pred'] loss['part_loss'] = add_part_loss(multi_class_labels, logits, part_score_weights) total_loss = loss['part_loss'] if dlc_cfg.intermediate_supervision: logits_intermediate = heads['part_loss_interm'] loss['part_loss_interm'] = add_part_loss(multi_class_labels, logits_intermediate, part_score_weights) total_loss = total_loss + loss['part_loss_interm'] if dlc_cfg.location_refinement: locref_pred = heads['locref'] locref_targets = batch[Batch.locref_targets] locref_weights = batch[Batch.locref_mask] loss_func = losses.huber_loss if dlc_cfg.locref_huber_loss else tf.losses.mean_squared_error loss['locref_loss'] = dlc_cfg.locref_loss_weight * loss_func( locref_targets, locref_pred, locref_weights) total_loss = total_loss + loss['locref_loss'] # loss['total_loss'] = slim.losses.get_total_loss(add_regularization_losses=params.regularize) loss['total_loss'] = total_loss #%% for k, t in loss.items(): TF.summary.scalar(k, t) TF.summary.merge_all() #%% # restore from snapshot if trainingsnapshot == 0: variables_to_restore = slim.get_variables_to_restore( include=["resnet_v1"]) else: variables_to_restore = slim.get_variables_to_restore() restorer = TF.train.Saver(variables_to_restore) #%% Init session config_TF = TF.ConfigProto() config_TF.gpu_options.allow_growth = True sess = TF.Session(config=config_TF) coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders) TF.summary.FileWriter(dlc_cfg.log_dir, sess.graph) learning_rate, train_op = get_optimizer(total_loss, dlc_cfg) #%% sess.run(TF.global_variables_initializer()) sess.run(TF.local_variables_initializer()) # Restore the one variable from disk restorer.restore(sess, dlc_cfg.init_weights) print('Restored variables from\n{}\n'.format(dlc_cfg.init_weights)) #%% lr_gen = LearningRate(dlc_cfg) #%% dlc_params_outdir = Path( dlc_cfg.init_weights).parent / 'dlc_params_mat' / '{}'.format( trainingsnapshot_name) if not os.path.isdir(dlc_params_outdir): assert Path(dlc_cfg.init_weights).parent os.makedirs(dlc_params_outdir) print(dlc_params_outdir) #%% biases = [ v for v in tf.global_variables() if v.name == "pose/part_pred/block4/biases:0" ][0] weights = [ v for v in tf.global_variables() if v.name == "pose/part_pred/block4/weights:0" ][0] if dlc_cfg.location_refinement: biases_locref = [ v for v in tf.global_variables() if v.name == "pose/locref_pred/block4/biases:0" ][0] weights_locref = [ v for v in tf.global_variables() if v.name == "pose/locref_pred/block4/weights:0" ][0] # locref_pred #%% current_lr = lr_gen.get_lr(0) if dlc_cfg.location_refinement: [_, biases_out, weights_out, bias_locref_out, weight_locref_out] = sess.run( [train_op, biases, weights, biases_locref, weights_locref], feed_dict={learning_rate: current_lr}) ss = os.path.join(dlc_params_outdir, 'dlc_params.mat') sio.savemat( ss, { 'weight': weights_out, 'bias': biases_out, 'weight_locref': weight_locref_out, 'bias_locref': bias_locref_out }) else: [_, biases_out, weights_out] = sess.run([train_op, biases, weights], feed_dict={learning_rate: current_lr}) ss = os.path.join(dlc_params_outdir, 'dlc_params.mat') sio.savemat(ss, {'weight': weights_out, 'bias': biases_out}) print('\nStored output in\n{}\n'.format(str(ss))) sess.close() coord.request_stop() return
def _compute_targets(self): from deeplabcut.pose_estimation_tensorflow.dataset.factory import \ create as create_dataset dlc_config = copy.deepcopy(self.dlc_config) dlc_config['deterministic'] = True # switch to default dataset_type to produce expected batch output dlc_config['dataset_type'] = 'default' dataset = create_dataset(dlc_config) nt = len(self.idxs['vis']['train']) # number of training frames # assert nt >= 1 nj = max([dat_.joints[0].shape[0] for dat_ in dataset.data]) stride = dlc_config['stride'] def extract_frame_num(img_path): return int(img_path.rsplit('/', 1)[-1][3:].split('.')[0]) frame_idxs = [] joinss = [] counter = 0 while counter < nt: data = dataset.next_batch() data_keys = list(data.keys()) # inputs = 0 # part_score_targets = 1 # part_score_weights = 2 # locref_targets = 3 # locref_mask = 4 # pairwise_targets = 5 # pairwise_mask = 6 # data_item = 7 im_path = data[data_keys[5]].im_path # skip if frame belongs to another video if im_path.find('/' + self.video_name + '/') == -1: continue # skip if image has already been processed frame_idx = extract_frame_num(im_path) if frame_idx in frame_idxs: continue # inputs = data[data_keys[0]].squeeze() # part_score_targets = data[data_keys[1]].squeeze() # multi class labels # part_score_weights = data[data_keys[2]].squeeze() # locref_targets = data[data_keys[3]].squeeze() # locref_mask = data[data_keys[4]].squeeze() data_item = data[data_keys[-1]] joinss += [np.copy(data_item.joints[0])] frame_idxs.append(frame_idx) counter += 1 # to find 2D coordinates, we must update targets_2d = np.zeros((nt, nj, 2)) * np.nan # nt x nj x 2 for ntt in range(nt): cjoin = joinss[ntt] # D x nj njtt = cjoin.shape[0] for njj_id in range(njtt): njj = cjoin[njj_id][0] joinss_ntt_njj = cjoin[njj_id][1:] targets_2d[ntt, njj] = np.flip( (joinss_ntt_njj - stride / 2) / stride) nx_out = int(data[data_keys[1]].squeeze().shape[0]) ny_out = int(data[data_keys[1]].squeeze().shape[1]) return targets_2d, frame_idxs, nx_out, ny_out
def get_targets(task, date, shuffle): from deeplabcut.pose_estimation_tensorflow.dataset.factory import ( create as create_dataset, ) #%% #task = 'reach' #date = '2020-02-19' #shuffle = 0 #%% data_info = DataLoader(task) # Load project configuration cfg = get_model_config(task, data_info.model_data_dir, scorer=data_info.scorer, date=date) # Load training configuration for cfg file dlc_cfg = get_train_config(cfg, shuffle=shuffle) # update _, _, dlc_cfg = load_dlc_snapshot(dlc_cfg) # %% dlc_cfg["deterministic"] = True #%% Create dataset dataset = create_dataset(dlc_cfg) nt = len(dataset.data) # number of training frames assert nt >= 1 nj = max([dat_.joints[0].shape[0] for dat_ in dataset.data]) ncolors, nxraw, nyraw = dataset.data[0].im_size #%% stride = dlc_cfg['stride'] #%% #clip = VideoFileClip(str(dlc_cfg.video_path)) #%% # TO DO: make exact should be around 1/8 for resnet 50 nx = nxraw // 4 ny = nyraw // 4 #%% extract_frame_num = lambda x: int(x.split("/")[-1].split(".")[0][3:]) frame_ids = [] # frame_imgs = [] counter = 0 #datas = [] datas = np.zeros((nt, nx, ny, nj)) joinss = [] #%% while counter < (nt): # for ii in range(nt): data = dataset.next_batch() data_keys = list(data.keys()) # inputs = 0 # part_score_targets = 1 # part_score_weights = 2 # locref_targets = 3 # locref_mask = 4 # pairwise_targets = 5 # pairwise_mask = 6 # data_item = 7 data_item = data[data_keys[-1]] joinss += [np.copy(data_item.joints[0])] inputs = data[data_keys[0]].squeeze() print(inputs.shape) part_score_targets = data[data_keys[1]].squeeze() # multi class labels nx, ny = part_score_targets.shape[0], part_score_targets.shape[1] # part_score_weights = data[data_keys[2]].squeeze() # locref_targets = data[data_keys[3]].squeeze() # locref_mask = data[data_keys[4]].squeeze() frame_id = extract_frame_num(data[data_keys[5]].im_path) # print(part_score_targets.max((0, 1))) print(frame_id) if frame_id in frame_ids: continue else: print("Adding frame {}".format(frame_id)) frame_ids.append(frame_id) datas[counter, :nx, :ny, :] = part_score_targets #datas.append(part_score_targets) counter += 1 #%% #datas = np.stack(datas, 0) datas = datas[:, :nx, :ny, :] #nx, ny = datas.shape[1:3] frame_ids = np.asarray(frame_ids) #%% # ignore when tongue is not present? # assert datas.sum((1, 2)).min() >= 1 print(datas.sum((1, 2)).min()) # %% # To find 2D coordinates, we must update target2d_train = np.zeros((nt, nj, 2)) * np.nan # nt x nj x 2 for ntt in range(nt): cjoin = joinss[ntt] # D x nj njtt = cjoin.shape[0] for njj_id in range(njtt): njj = cjoin[njj_id][0] joinss_ntt_njj = cjoin[njj_id][1:] target2d_train[ntt, njj] = np.flip( (joinss_ntt_njj - stride / 2) / stride) # %% # visualize some randomly # ntt = 23 # njj = 0 # center_loc = target2d_train[ntt, njj] # plt.imshow(datas[ntt, :, :, njj]) # plt.plot(center_loc[1], center_loc[0], 'ro') # plt.show(block=True) # %% # target is shared for all snapshots regargless of whathever else nx_out = int(part_score_targets.shape[0]) ny_out = int(part_score_targets.shape[1]) return frame_ids, target2d_train, nx_out, ny_out, datas
def train(config_yaml, displayiters, saveiters, maxiters, max_to_keep=5, projection_matrices=None, multiview_step=None, snapshot_index=None): start_path = os.getcwd() os.chdir(str(Path(config_yaml).parents[0]) ) #switch to folder of config_yaml (for logging) setup_logging() cfg = load_config(config_yaml) cfg['batch_size'] = 1 #in case this was edited for analysis. cfg['projection_matrices'] = projection_matrices cfg['multiview_step'] = multiview_step # at this step, jittering the image sizes won't help # also, if we jitter the sizes then we would have to undo the jitter before projecting to 3D, so we may as well keep the image size constant if multiview_step == 2: cfg.global_scale = 1.0 cfg.scale_jitter_lo = 1.0 cfg.scale_jitter_up = 1.0 # also found best results with this optimizer and lr print('switching to hardcoded Adam optimizer for this step') cfg.optimizer = 'adam' cfg.adam_lr = 0.0001 dataset = create_dataset(cfg) batch_spec = get_batch_spec(cfg) batch, enqueue_op, placeholders = setup_preloading(batch_spec) losses = pose_net(cfg).train(batch) total_loss = losses['total_loss'] for k, t in losses.items(): tf.summary.scalar(k, t) merged_summaries = tf.summary.merge_all() if snapshot_index is None: variables_to_restore = slim.get_variables_to_restore( include=["resnet_v1"]) else: variables_to_restore = slim.get_variables_to_restore(exclude=[ op.name for op in tf.global_variables(scope='.*reweighting.*') ]) cfg.init_weights = os.path.join(os.path.dirname(config_yaml), 'snapshot-%d' % snapshot_index) restorer = tf.train.Saver(variables_to_restore) saver = tf.train.Saver( max_to_keep=max_to_keep ) # selects how many snapshots are stored, see https://github.com/AlexEMG/DeepLabCut/issues/8#issuecomment-387404835 sess = tf.Session() coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders) train_writer = tf.summary.FileWriter(cfg.log_dir, sess.graph) learning_rate, train_op = get_optimizer(total_loss, cfg) sess.run(tf.global_variables_initializer()) sess.run(tf.local_variables_initializer()) # Restore variables from disk. restorer.restore(sess, cfg.init_weights) if maxiters == None: max_iter = int(cfg.multi_step[-1][1]) else: max_iter = min(int(cfg.multi_step[-1][1]), int(maxiters)) #display_iters = max(1,int(displayiters)) print("Max_iters overwritten as", max_iter) if displayiters == None: display_iters = max(1, int(cfg.display_iters)) else: display_iters = max(1, int(displayiters)) print("Display_iters overwritten as", display_iters) if saveiters == None: save_iters = max(1, int(cfg.save_iters)) else: save_iters = max(1, int(saveiters)) print("Save_iters overwritten as", save_iters) cum_loss = 0.0 lr_gen = LearningRate(cfg) stats_path = Path(config_yaml).with_name('learning_stats.csv') lrf = open(str(stats_path), 'w') print("Training parameter:") print(cfg) print("Starting training....") for it in range(max_iter + 1): current_lr = lr_gen.get_lr(it) [_, loss_val, summary] = sess.run([train_op, total_loss, merged_summaries], feed_dict={learning_rate: current_lr}) cum_loss += loss_val train_writer.add_summary(summary, it) if it % display_iters == 0 and it > 0: average_loss = cum_loss / display_iters cum_loss = 0.0 logging.info("iteration: {} loss: {} lr: {}".format( it, "{0:.4f}".format(average_loss), current_lr)) lrf.write("{}, {:.5f}, {}\n".format(it, average_loss, current_lr)) lrf.flush() # Save snapshot if (it % save_iters == 0 and it != 0) or it == max_iter: model_name = cfg.snapshot_prefix saver.save(sess, model_name, global_step=it) lrf.close() sess.close() coord.request_stop() coord.join([thread]) #return to original path. os.chdir(str(start_path))