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))
