def train_autoencoder(model_path,
train_data_path,
image_size,
report_rate=100,
learning_rate=1e-4,
num_epoch=50,
batch_size=10,
capacity=3000,
min_after_dequeue=800):
from data.common import TfReader
with tf.Graph().as_default():
# Read training data.
train_data = TfReader(data_path=train_data_path,
size=(image_size, image_size),
num_epochs=num_epoch)
images, classes = train_data.read(batch_size=batch_size,
capacity=capacity,
min_after_dequeue=min_after_dequeue)
y, z = build_autoencoder(images)
# Cost function measures pixel-wise difference
cost = tf.reduce_sum(tf.square(y - images))
optimize(cost=cost,
save_path=model_path,
report_rate=report_rate,
scope=None,
learning_rate=learning_rate)
def test_autoencoder(model_path,
test_data_path,
image_size,
report_rate=100,
batch_size=1,
capacity=3000,
min_after_dequeue=800):
from PIL import Image
import numpy as np
import os
from data.common import TfReader
with tf.Graph().as_default():
# Read test data.
test_data = TfReader(data_path=test_data_path,
size=(image_size, image_size))
images, classes = test_data.read(batch_size=batch_size,
capacity=capacity,
min_after_dequeue=min_after_dequeue)
y, z = build_autoencoder(images)
# Cost function measures pixel-wise difference
cost = tf.reduce_sum(tf.square(y - images))
# Run session.
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
saver = tf.train.Saver()
with open(os.path.join(model_path, "checkpoint")) as checkpoint_file:
checkpoint_name = checkpoint_file.readline().strip().split(
':', 1)[1].strip()[1:-1]
print "Using model: " + checkpoint_name
with tf.Session() as sess:
sess.run(init_op)
saver.restore(sess, os.path.join(model_path, checkpoint_name))
tf.train.start_queue_runners(sess)
overall_cost = 0.0
step_count = 0.0
try:
while True:
original, decoded, predictions, current_cost = sess.run(
[images, y, z, cost])
step_count += 1.0
overall_cost += current_cost
if step_count % report_rate == 0:
Image.fromarray(np.uint8(decoded[0]), "RGB") \
.save("decoded_step_{0}.jpg".format(step_count))
Image.fromarray(np.uint8(original[0]), "RGB") \
.save("original_step_{0}.jpg".format(step_count))
print "{0} steps passed with cost of {1}/{2}." \
.format(step_count, current_cost, overall_cost / step_count)
except tf.errors.OutOfRangeError:
print "All images used in {0} steps.".format(step_count)
finally:
print "Final cost: {0}.".format(overall_cost / step_count)
def extract_image(input_path,
resize,
limit,
output_path,
regression,
report_rate=50):
from data.common import TfReader
from PIL import Image
import numpy as np
import tensorflow as tf
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
reader = TfReader(data_path=input_path,
regression=regression,
size=(resize, resize))
img_op, label_op = reader.read(batch_size=1)
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
with tf.Session() as sess:
sess.run([init_op])
tf.train.start_queue_runners(sess)
step_count = 0
label_wrote_count = {}
try:
while True:
step_count += 1
img, label = sess.run([img_op, label_op])
img = img[0]
label = label[0]
output_dir = os.path.join(output_path, str(label))
# Count name wrote times.
if label not in label_wrote_count:
label_wrote_count[label] = 0
if not os.path.isdir(output_dir):
os.mkdir(output_dir)
# Skip extra images.
elif limit and label_wrote_count[label] >= limit:
continue
Image.fromarray(np.uint8(img), "RGB").save(
os.path.join(output_dir,
"step_{0}.jpg".format(step_count)))
label_wrote_count[label] += 1
if step_count % report_rate == 0:
print "{0} steps passed with label wrote: ".format(
step_count),
print label_wrote_count
except tf.errors.OutOfRangeError:
print "All images used in {0} steps.".format(step_count)
def train(model_path,
train_data_path,
class_count,
image_size,
image_channel=3,
report_rate=100,
build=build_cnn,
regression=False,
scope=None,
learning_rate=1e-4,
num_epoch=50,
batch_size=10,
capacity=3000,
min_after_dequeue=800):
from data.common import TfReader
with tf.Graph().as_default():
# Read training data.
train_data = TfReader(data_path=train_data_path,
regression=regression,
size=(image_size, image_size),
num_epochs=num_epoch)
images, classes = train_data.read(batch_size=batch_size,
capacity=capacity,
min_after_dequeue=min_after_dequeue)
y, y_pred_cls = build(input_tensor=images,
num_class=class_count,
image_size=image_size,
image_channel=image_channel)
if regression:
cost = tf.reduce_sum(tf.pow(tf.transpose(y) - classes,
2)) / (2 * batch_size)
else:
# Calculate cross-entropy for each image.
# This function calculates the softmax internally, so use the output layer directly.
# The input label is of type int in [0, num_class).
cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=y, labels=classes)
# Calculate average cost across all images.
cost = tf.reduce_mean(cross_entropy)
optimize(cost=cost,
save_path=model_path,
report_rate=report_rate,
scope=scope,
learning_rate=learning_rate)
def test(config, alex):
from io import BytesIO
from PIL import Image
from zipfile import ZipFile
import time
try:
os.mkdir(os.path.join(config["save_root"], "crimtane"))
except OSError:
pass
print("==> test started at {0} for {1}.".format(
datetime.now().strftime("%Y-%m-%d %H:%M"), alex))
print("---- CONFIG DUMP ----")
print(json.dumps(config, indent=1))
print("---- END ----")
print("--> building models...")
feature_module = alex()
gender_module = NnGender(feature=feature_module.feature)
image, label = TfReader(data_path=config["test_data"]["path"], size=(256, 256),
num_epochs=config["test_data"]["epoch"]) \
.read(batch_size=config["test_data"]["batch_size"])
statistics = ConfusionMatrix(2)
# MonitoredTrainingSession takes care of session initialization,
# restoring from a checkpoint, saving to a checkpoint, and closing when done
# or an error occurs.
print("--> starting session...")
if config["keep_zip"] > 0:
zip_file = ZipFile(
os.path.join(config["save_root"], "crimtane",
"test_result_{0}.zip".format(int(time.time()))), 'w')
else:
zip_file = DummyFile()
hooks = [
LoadInitialValueHook(module_list=[feature_module, gender_module],
save_path=config["save_root"])
]
with tf.train.MonitoredTrainingSession(hooks=hooks) as mon_sess:
accumulated_accuracy = 0
test_step = 0
file_count = 0
try:
while not mon_sess.should_stop():
test_step += 1
image_batch, label_batch = mon_sess.run([image, label])
prediction_value = mon_sess.run(
gender_module.prediction,
feed_dict={feature_module.image_input: image_batch})
accuracy_value = 1 - np.mean(
np.abs(np.transpose(prediction_value) - label_batch))
accumulated_accuracy += accuracy_value
statistics.update(predictions=prediction_value,
truth=label_batch)
if test_step % config["report_rate"] == 0:
print(" * step ({0}) accuracy: {1:8}".format(
test_step, accumulated_accuracy / test_step))
if config["keep_zip"] > 0:
for image_bytes, label_float, prediction_float in zip(
image_batch, label_batch, prediction_value):
prediction_string = str(
int(prediction_float +
(0.5 if prediction_float > 0 else -0.5)))
label_string = str(int(label_float))
jpeg_bytes = BytesIO()
Image.fromarray(np.uint8(image_bytes),
"RGB").save(jpeg_bytes, format="JPEG")
zip_file.writestr(
"{0}/{1}_{2:05}.jpg".format(
label_string, prediction_string, file_count),
jpeg_bytes.getvalue())
file_count += 1
except tf.errors.OutOfRangeError as e:
print("no more data: {0}".format(repr(e)))
except KeyboardInterrupt as e:
print("\ncanceled: {0}".format(repr(e)))
zip_file.close()
with open(os.path.join(config["save_root"], "gan_vgg.log"),
'a') as log_file:
message = "==> test for {2} completed at {0} in {1} steps.".format(
datetime.now().strftime("%Y-%m-%d %H:%M"), test_step, alex)
log_file.write(message + "\n")
print(message)
log_file.write("---- CONFIG DUMP ----\n")
json.dump(config, log_file, indent=1)
log_file.write("\n---- END ----\n")
if config["keep_zip"] > 0:
keep_zip = config["keep_zip"]
import glob
zips = glob.glob(
os.path.join(config["save_root"], "crimtane",
"test_result_*.zip"))
for a_zip in sorted(zips, reverse=True):
if keep_zip > 0:
keep_zip -= 1
continue
os.remove(a_zip)
message = "result wrote to {0}.".format(zip_file.filename)
log_file.write(message + "\n")
print(message)
message = "overall result: {0:8}".format(accumulated_accuracy /
test_step)
log_file.write(message + "\n")
print(message)
message = statistics.generate_result()
log_file.write(message + "\n")
print(message)
def adaption(config):
try:
os.mkdir(os.path.join(config["save_root"], "adamantite"))
except OSError:
pass
print("==> adaption started at {0}.".format(
datetime.now().strftime("%Y-%m-%d %H:%M")))
print("---- CONFIG DUMP ----")
print(json.dumps(config, indent=1))
print("---- END ----")
print("--> building models...")
source_feature_module = Source()
target_feature_module = Target()
target_feature_module.override_saver_for_init_by(
source_model=source_feature_module)
discriminator_module = NnClassification(
feature=target_feature_module.feature)
source_image, _ = TfReader(data_path=config["source_data"]["path"], size=(256, 256),
num_epochs=config["source_data"]["epoch"]) \
.read(batch_size=config["source_data"]["batch_size"])
target_image, _ = TfReader(data_path=config["target_data"]["path"], size=(256, 256),
num_epochs=config["target_data"]["epoch"]) \
.read(batch_size=config["target_data"]["batch_size"])
global_step_op = tf.Variable(0, trainable=False, name="global_step")
var_d = tf.get_collection(key=tf.GraphKeys.TRAINABLE_VARIABLES,
scope=discriminator_module.variable_scope)
optimizer_d = tf.train.AdamOptimizer(learning_rate=config["alternative_learning_rate"]) \
.minimize(loss=discriminator_module.loss,
global_step=global_step_op,
var_list=var_d,
colocate_gradients_with_ops=True)
optimizer_m = tf.train.AdamOptimizer(learning_rate=config["learning_rate"]) \
.minimize(loss=discriminator_module.loss,
global_step=global_step_op,
var_list=target_feature_module.trainable_list,
colocate_gradients_with_ops=True)
print("optimizer_d variables:", end='')
for index, var in enumerate(var_d):
if index % 2 == 0:
print("\n", end='')
print(" {0}".format(var), end='')
print("\n", end='')
print("optimizer_m variables:", end='')
for index, var in enumerate(target_feature_module.trainable_list):
if index % 2 == 0:
print("\n", end='')
print(" {0}".format(var), end='')
print("\n", end='')
print("--> starting session...")
if config["checkpointing"]:
checkpoint = os.path.join(config["save_root"], "adamantite")
else:
checkpoint = None
# no need to init when checkpoint exist
if checkpoint and os.path.exists(os.path.join(checkpoint, "checkpoint")):
hooks = [
EndSavingHook(
module_list=[target_feature_module, discriminator_module],
save_path=config["save_root"])
]
else:
hooks = [
EndSavingHook(
module_list=[target_feature_module, discriminator_module],
save_path=config["save_root"]),
LoadInitialValueHook(
module_list=[source_feature_module, target_feature_module],
save_path=config["save_root"])
]
with tf.train.MonitoredTrainingSession(
hooks=hooks, checkpoint_dir=checkpoint) as mon_sess:
global_step = -1
cost_d = -1
cost_m = -1
accuracy_d = -1
step_for_report = 0
step_for_header = 0
try:
while not mon_sess.should_stop():
# read image and compute the feature
source_image_batch, target_image_batch = mon_sess.run(
[source_image, target_image])
source_feature_batch, target_feature_batch = mon_sess.run(
[
source_feature_module.feature,
target_feature_module.feature
],
feed_dict={
source_feature_module.image_input: source_image_batch,
target_feature_module.image_input: target_image_batch
})
# discriminator
if "discriminator" in config["adaption_mode"]:
import random
combined = list(
zip(
np.concatenate(
(source_feature_batch, target_feature_batch),
axis=0),
[1] * config["source_data"]["batch_size"] +
[0] * config["target_data"]["batch_size"]))
random.shuffle(combined)
features, labels = zip(*combined)
_, global_step, cost_d = mon_sess.run(
[
optimizer_d, global_step_op,
discriminator_module.loss
],
feed_dict={
target_feature_module.feature: features,
discriminator_module.label_input: labels
})
# generator
if "generator" in config["adaption_mode"]:
epoch_multiplier_d = 1
cost_m = 0
accumulated_cost = 0
for _ in range(epoch_multiplier_d):
_, global_step, current_cost = mon_sess.run(
[
optimizer_m, global_step_op,
discriminator_module.loss
],
feed_dict={
target_feature_module.image_input:
target_image_batch,
discriminator_module.label_input:
[1] * config["target_data"]["batch_size"]
})
accumulated_cost += current_cost
cost_m = accumulated_cost / epoch_multiplier_d
# determine accuracy
features = np.concatenate(
(source_feature_batch, target_feature_batch), axis=0)
labels = [1] * config["source_data"]["batch_size"] + [
0
] * config["target_data"]["batch_size"]
prediction_batch = mon_sess.run(
discriminator_module.prediction,
feed_dict={target_feature_module.feature: features})
accuracy_d = 1 - np.mean(
np.abs(np.transpose(prediction_batch) - labels))
# report progress
if global_step >= step_for_header:
step_for_header = global_step + config["report_rate"] * 25
print(
"\n step cost_m cost_d accuracy\n - - - - - - - - - - - - - - -"
)
if global_step >= step_for_report:
step_for_report = global_step + config["report_rate"]
print(" * {0:8} {1:10.4f} {2:10.4f} {3:8.4f}%".format(
global_step, cost_m, cost_d, accuracy_d * 100))
except tf.errors.OutOfRangeError as e:
print("no more data: {0}".format(repr(e)))
except KeyboardInterrupt as e:
print("\ncanceled: {0}".format(repr(e)))
with open(os.path.join(config["save_root"], "gan_vgg.log"),
'a') as log_file:
message = "==> adaption completed at {0} in {1} steps.".format(
datetime.now().strftime("%Y-%m-%d %H:%M"), global_step)
log_file.write(message + "\n")
print(message)
log_file.write("---- CONFIG DUMP ----\n")
json.dump(config, log_file, indent=1)
log_file.write("\n---- END ----\n")
message = "cost_d: {0:8}, cost_m: {1:8}, accuracy: {2:8}".format(
cost_d, cost_m, accuracy_d)
log_file.write(message + "\n")
print(message)
os.remove(os.path.join(config["save_root"], "checkpoint"))
def pre_train(config):
try:
os.mkdir(os.path.join(config["save_root"], "adamantite"))
except OSError:
pass
print("==> pre-train started at {0}.".format(
datetime.now().strftime("%Y-%m-%d %H:%M")))
print("---- CONFIG DUMP ----")
print(json.dumps(config, indent=1))
print("---- END ----")
print("--> building models...")
source_feature_module = Source()
gender_module = NnGender(feature=source_feature_module.feature)
image, label = TfReader(data_path=config["source_data"]["path"], size=(256, 256),
num_epochs=config["source_data"]["epoch"]) \
.read(batch_size=config["source_data"]["batch_size"])
global_step_op = tf.Variable(0, trainable=False, name="global_step")
var_to_train = tf.get_collection(
tf.GraphKeys.TRAINABLE_VARIABLES,
gender_module.variable_scope) + source_feature_module.trainable_list
optimizer = tf.train.AdamOptimizer(
learning_rate=config["learning_rate"]).minimize(
loss=gender_module.loss,
global_step=global_step_op,
var_list=var_to_train,
colocate_gradients_with_ops=True)
print("optimizer variables:", end='')
for index, var in enumerate(var_to_train):
if index % 2 == 0:
print("\n", end='')
print(" {0}".format(var), end='')
print("\n", end='')
print("--> starting session...")
if config["checkpointing"]:
checkpoint = os.path.join(config["save_root"], "adamantite")
else:
checkpoint = None
hooks = [
EndSavingHook(module_list=[source_feature_module, gender_module],
save_path=config["save_root"])
]
with tf.train.MonitoredTrainingSession(
hooks=hooks, checkpoint_dir=checkpoint) as mon_sess:
global_step = -1
current_cost = -1
try:
while not mon_sess.should_stop():
image_batch, label_batch = mon_sess.run([image, label])
_, global_step, current_cost = mon_sess.run(
[optimizer, global_step_op, gender_module.loss],
feed_dict={
source_feature_module.image_input: image_batch,
gender_module.label_input: label_batch
})
if global_step % config["report_rate"] == 0:
print(" * step ({0}) cost: {1:8}".format(
global_step, current_cost))
except tf.errors.OutOfRangeError as e:
print("no more data: {0}".format(repr(e)))
except KeyboardInterrupt as e:
print("\ncanceled: {0}".format(repr(e)))
with open(os.path.join(config["save_root"], "gan_vgg.log"),
'a') as log_file:
message = "==> pre-train completed at {0} in {1} steps.".format(
datetime.now().strftime("%Y-%m-%d %H:%M"), global_step)
log_file.write(message + "\n")
print(message)
log_file.write("---- CONFIG DUMP ----\n")
json.dump(config, log_file, indent=1)
log_file.write("\n---- END ----\n")
message = "cost: {0:8}".format(current_cost)
log_file.write(message + "\n")
print(message)
os.remove(os.path.join(config["save_root"], "checkpoint"))
def test(model_path,
test_data_path,
class_count,
image_size,
image_channel=3,
report_rate=100,
build=build_cnn,
regression=False,
batch_size=1,
capacity=3000,
min_after_dequeue=800):
from data.common import TfReader
import os
with tf.Graph().as_default():
# Read test data.
train_data = TfReader(data_path=test_data_path,
regression=regression,
size=(image_size, image_size))
images, classes = train_data.read(batch_size=batch_size,
capacity=capacity,
min_after_dequeue=min_after_dequeue)
y, y_pred_cls = build(input_tensor=images,
num_class=class_count,
image_size=image_size,
image_channel=image_channel)
if regression:
prediction_op = tf.squeeze(y)
correct_prediction = tf.abs(tf.transpose(y) - classes)
statistics = RegressionBias()
else:
prediction_op = y_pred_cls
correct_prediction = tf.equal(y_pred_cls, classes)
statistics = ConfusionMatrix(class_count=class_count)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
# Run session.
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
saver = tf.train.Saver()
with open(os.path.join(model_path, "checkpoint")) as checkpoint_file:
checkpoint_name = checkpoint_file.readline().strip().split(
':', 1)[1].strip()[1:-1]
print "Using model: " + checkpoint_name
with tf.Session() as sess:
sess.run(init_op)
saver.restore(sess, os.path.join(model_path, checkpoint_name))
tf.train.start_queue_runners(sess)
overall_accuracy = 0.0
step_count = 0.0
try:
while True:
predictions, truth, current_accuracy = sess.run(
[prediction_op, classes, accuracy])
step_count += 1.0
overall_accuracy += current_accuracy
statistics.update(predictions=predictions, truth=truth)
if step_count % report_rate == 0:
print "{0} steps passed with result of {1}/{2}."\
.format(step_count, current_accuracy, overall_accuracy / step_count)
print "Sample predictions: {0}".format(predictions)
except tf.errors.OutOfRangeError:
print "All images used in {0} steps.".format(step_count)
finally:
print "Final result: {0}.".format(overall_accuracy /
step_count)
statistics.print_result()
return overall_accuracy / step_count