def read_audio_files(audiodir, parsed_features):
wav1 = tf.read_file(
tf.string_join([audiodir, parsed_features['comb/file1']],
separator='/'))
wav1 = tf.contrib.ffmpeg.decode_audio(wav1,
file_format='wav',
samples_per_second=44100,
channel_count=1)
wav1 = tf.concat(
[tf.zeros([parsed_features['comb/sig1_sample_delay'], 1]), wav1],
axis=0)
wav2 = tf.read_file(
tf.string_join([audiodir, parsed_features['comb/file2']],
separator='/'))
wav2 = tf.contrib.ffmpeg.decode_audio(wav2,
file_format='wav',
samples_per_second=44100,
channel_count=1)
wav2 = tf.concat(
[tf.zeros([parsed_features['comb/sig2_sample_delay'], 1]), wav2],
axis=0)
parsed_features['sig1/samples'] = tf.divide(wav1,
1.5 * tf.reduce_max(wav1))
parsed_features['sig2/samples'] = tf.divide(wav2,
1.5 * tf.reduce_max(wav2))
return parsed_features
def preprocessing_fn(inputs):
"""Preprocess input columns into transformed columns."""
outputs = {}
for key in numerical_feats:
outputs[key] = tf.cast(tft.bucketize(inputs[key], 20),
tf.float32) / 20.0 - 0.5
outputs["campaignCost_mod"] = inputs["campaignCost"] / 100.0
inputs["game_zone"] = tf.string_join(
[inputs["sourceGameId"], inputs["zone"]], separator="_")
inputs["game_campaignId"] = tf.string_join(
[inputs["sourceGameId"], inputs["campaignId"]], separator="_")
for key in categorical_feats + ["game_zone", "game_campaignId"]:
vocab = tft.vocabulary(inputs[key],
vocab_filename=key,
frequency_threshold=100)
outputs[key] = tft.apply_vocabulary(inputs[key],
vocab,
default_value=0)
outputs["label"] = inputs["label"]
outputs["key"] = inputs["key"]
return outputs
def _load_and_decode(data_location, im_la_files):
"""
Loads and decodes images and labels from a location specified as a string
Args:
data_location: a String or tf.string with the location of the Apolloscape dataset
im_la_files: a tf.string with the location of the image and the label, separated by a tab.
Returns:
"""
data_location = tf.cast(data_location, tf.string)
im_la_files = tf.cast(im_la_files, tf.string)
im_la_files_split = tf.string_split([im_la_files], '\t')
im_file = im_la_files_split.values[0]
la_file = im_la_files_split.values[1]
im_string = tf.read_file(tf.string_join([data_location, im_file]))
im_dec = tf.image.decode_jpeg(im_string)
la_string = tf.read_file(tf.string_join([data_location, la_file]))
la_dec = tf.image.decode_png(la_string)[..., 0]
return im_dec, la_dec, im_file, la_file
def __init__(self, data_path, filenames_file, params, dataset, mode):
self.data_path = data_path
self.params = params
self.dataset = dataset
self.mode = mode
self.left_image_batch = None
self.right_image_batch = None
input_queue = tf.train.string_input_producer([filenames_file],
shuffle=False)
line_reader = tf.TextLineReader()
_, line = line_reader.read(input_queue)
split_line = tf.string_split([line]).values
# we load only one image for test, except if we trained a stereo model
left_image_path = tf.string_join([self.data_path, split_line[0]])
right_image_path = tf.string_join([self.data_path, split_line[1]])
left_image_o = self.read_image(left_image_path)
right_image_o = self.read_image(right_image_path)
if mode == 'train':
# randomly flip images
do_flip = tf.random_uniform([], 0, 1)
left_image = left_image_o #tf.cond(do_flip > 0.5, lambda: tf.image.flip_left_right(right_image_o), lambda: left_image_o)
right_image = right_image_o #tf.cond(do_flip > 0.5, lambda: tf.image.flip_left_right(left_image_o), lambda: right_image_o)
# randomly augment images
do_augment = tf.random_uniform([], 0, 1)
left_image, right_image = tf.cond(
do_augment > 0.5,
lambda: self.augment_image_pair(left_image, right_image),
lambda: (left_image, right_image))
left_image.set_shape([None, None, 3])
right_image.set_shape([None, None, 3])
# capacity = min_after_dequeue + (num_threads + a small safety margin) * batch_size
min_after_dequeue = 2048
capacity = min_after_dequeue + 4 * params.batch_size
self.left_image_batch, self.right_image_batch = tf.train.shuffle_batch(
[left_image, right_image], params.batch_size, capacity,
min_after_dequeue, params.num_threads)
elif mode == 'test':
print('mode: ', mode)
left_image = left_image_o
right_image = right_image_o
left_image.set_shape([None, None, 3])
right_image.set_shape([None, None, 3])
# capacity = min_after_dequeue + (num_threads + a small safety margin) * batch_size
min_after_dequeue = 2048
capacity = min_after_dequeue + 4 * params.batch_size
self.left_image_batch, self.right_image_batch = tf.train.batch(
[left_image, right_image],
params.batch_size,
capacity=capacity,
num_threads=params.num_threads)
def log_text(F, G, params):
lookup_table = construct_vocab_lookup_table(params.vocab)
X_vocab = tf.expand_dims(tf.range(params.vocab_size), axis=0)
if params.use_embeddings:
X = embed_inputs(X_vocab, params, reuse=True)
else:
X = tf.one_hot(X_vocab, depth=params.vocab_size)
X_map_distribution = F(X, params.F, params)
X_map_indices = tf.argmax(X_map_distribution, axis=-1)
# X_vocab = tf.Print(X_vocab, [X_vocab], message="X_vocab", summarize=10)
# X_map_indices = tf.Print(
# X_map_indices, [X_map_indices], message="X_map_indices", summarize=10)
X_map_text = lookup_table.lookup(tf.to_int64(X_map_indices))
X_vocab_text = lookup_table.lookup(tf.to_int64(X_vocab))
X_text = tf.string_join([X_vocab_text, "->", X_map_text])
tf.summary.text("F_map", X_text)
Y_vocab = tf.expand_dims(tf.range(params.vocab_size), axis=0)
if params.use_embeddings:
Y = embed_inputs(Y_vocab, params, reuse=True)
else:
Y = tf.one_hot(Y_vocab, depth=params.vocab_size)
Y_map_distribution = G(Y, params.G, params)
Y_map_indices = tf.argmax(Y_map_distribution, axis=-1)
# Y_vocab = tf.Print(Y_vocab, [Y_vocab], message="Y_vocab", summarize=10)
# Y_map_indices = tf.Print(
# Y_map_indices, [Y_map_indices], message="Y_map_indices", summarize=10)
Y_map_text = lookup_table.lookup(tf.to_int64(Y_map_indices))
Y_vocab_text = lookup_table.lookup(tf.to_int64(Y_vocab))
Y_text = tf.string_join([Y_vocab_text, "->", Y_map_text])
tf.summary.text("G_map", Y_text)
def extract_data_from_path(self,path):
'''
Read each batch of data. Because the stored data is in h5py format, you need to support the tf.py_func function.
When reading data at the same time, pay attention to whether it is in the test stage.
'''
audio_filename = path
feature_filename = tf.string_join([tf.string_split([audio_filename],".").values[0],'.40logfbank'])
#Perform feature reading
#Read features in h5py files, while eliminating redundant dimensions
audio_feature = tf.squeeze(tf.py_func(self.read_feature,[feature_filename],[tf.float32]))
#Convert the read feature into a tensor, and the length of the feature should be recorded
audio_feature = tf.convert_to_tensor(audio_feature)
audio_length = tf.shape(audio_feature)[0]
if self.is_testing:
return {'audios':audio_feature,'audio_lengths':audio_length}
else:
#Read the target
label_filename = tf.string_join([tf.string_split([audio_filename],".").values[0],'.label'])
#Read the label in the h5py file, and you need to eliminate the extra dimensions
target_label = tf.squeeze(tf.py_func(self.read_label,[label_filename],[tf.int32]))
train_label = tf.squeeze(tf.py_func(self.read_train_label,[label_filename],[tf.int32]))
#Convert the read label to a tensor and record the length of the label
target_label = tf.convert_to_tensor(target_label)
target_length = tf.shape(target_label)[0]
train_label = tf.convert_to_tensor(train_label)
return {'audios':audio_feature,'audio_lengths':audio_length,'train_label':train_label,'target_label':target_label,'target_length':target_length}
def __init__(self, dataset, left_dir, right_dir, disp_dir):
self.dataset = dataset
self.left_dir = left_dir
self.right_dir = right_dir
self.disp_dir = disp_dir
self.left = None
self.right = None
self.disp = None
input_queue = tf.train.string_input_producer([self.dataset],
shuffle=False)
line_reader = tf.TextLineReader()
_, line = line_reader.read(input_queue)
split_line = tf.string_split([line], '.').values
self.left = tf.stack([
tf.cast(
self.read_image(tf.string_join([self.left_dir, line]),
[None, None, 3]), tf.float32)
], 0)
self.right = tf.stack([
tf.cast(
self.read_image(tf.string_join([self.right_dir, line]),
[None, None, 3]), tf.float32)
], 0)
self.disp = tf.stack([
tf.cast(
self.read_image(tf.string_join(
[self.disp_dir, split_line[0], '.png']), [None, None, 1],
dtype=tf.uint16), tf.float32)
], 0) / 256.
self.filename = split_line[0]
def read_and_decode(self, filename_queue):
img0_name = tf.string_join([self.img_dir, '/', filename_queue[0]])
img1_name = tf.string_join([self.img_dir, '/', filename_queue[1]])
img2_name = tf.string_join([self.img_dir, '/', filename_queue[2]])
img0 = tf.image.decode_png(tf.read_file(img0_name), channels=3)
img0 = tf.cast(img0, tf.float32)
img1 = tf.image.decode_png(tf.read_file(img1_name), channels=3)
img1 = tf.cast(img1, tf.float32)
img2 = tf.image.decode_png(tf.read_file(img2_name), channels=3)
img2 = tf.cast(img2, tf.float32)
flow_mask_noc_name = tf.string_join(
[self.img_dir, '/', filename_queue[3]])
flow_mask_occ_name = tf.string_join(
[self.img_dir, '/', filename_queue[4]])
flow_mask_noc = tf.image.decode_png(tf.read_file(flow_mask_noc_name),
dtype=tf.uint16,
channels=3)
flow_mask_noc = tf.cast(flow_mask_noc, tf.float32)
flow_mask_occ = tf.image.decode_png(tf.read_file(flow_mask_occ_name),
dtype=tf.uint16,
channels=3)
flow_mask_occ = tf.cast(flow_mask_occ, tf.float32)
flow_noc, mask_noc = self.extract_flow_and_mask(flow_mask_noc)
flow_occ, mask_occ = self.extract_flow_and_mask(flow_mask_occ)
return img0, img1, img2, flow_noc, flow_occ, mask_noc, mask_occ
def _parse_and_store_boxes(filename, dataset_directory, orig_dims):
filename_split = tf.unstack(tf.string_split([filename], "_").values[:-1],
num=3)
strip_filename = tf.string_join(filename_split, "_")
txt_dir = tf.cast(os.path.join(dataset_directory, 'panoptic_txt_weights/'),
tf.string)
txt_ext = tf.cast('_gtFine_instanceIds.txt', tf.string)
txt_filename = tf.string_join([txt_dir, strip_filename, txt_ext])
la_in_txt = tf.read_file(txt_filename)
la_in_txt = tf.string_split([la_in_txt], delimiter='\n').values
la_in_txt = tf.string_split(la_in_txt, delimiter=' ').values
la_in_int = tf.reshape(tf.string_to_number(la_in_txt, out_type=tf.int32),
[-1, 7])
# i_ids = la_in_int[:, 0]
weights = la_in_int[:, 6]
boxes_orig = la_in_int[:, 2:6]
boxes_format = convert_input_box_format(boxes_orig)
boxes_norm = normalize_boxes(boxes_format,
orig_height=orig_dims[0],
orig_width=orig_dims[1])
classes = la_in_int[:, 1]
instance_ids = la_in_int[:, 0]
return boxes_norm, classes, weights, instance_ids
def read_and_decode_distillation(self, filename_queue):
img1_name = tf.string_join([self.img_dir, '/', filename_queue[0]])
img2_name = tf.string_join([self.img_dir, '/', filename_queue[1]])
img1 = tf.image.decode_png(tf.read_file(img1_name), channels=3)
img1 = tf.cast(img1, tf.float32)
img2 = tf.image.decode_png(tf.read_file(img2_name), channels=3)
img2 = tf.cast(img2, tf.float32)
flow_occ_fw_name = tf.string_join([
self.fake_flow_occ_dir, '/flow_occ_fw_', filename_queue[2], '.png'
])
flow_occ_bw_name = tf.string_join([
self.fake_flow_occ_dir, '/flow_occ_bw_', filename_queue[2], '.png'
])
flow_occ_fw = tf.image.decode_png(tf.read_file(flow_occ_fw_name),
dtype=tf.uint16,
channels=3)
flow_occ_fw = tf.cast(flow_occ_fw, tf.float32)
flow_occ_bw = tf.image.decode_png(tf.read_file(flow_occ_bw_name),
dtype=tf.uint16,
channels=3)
flow_occ_bw = tf.cast(flow_occ_bw, tf.float32)
flow_fw, occ_fw = self.extract_flow_and_mask(flow_occ_fw)
flow_bw, occ_bw = self.extract_flow_and_mask(flow_occ_bw)
return img1, img2, flow_fw, flow_bw, occ_fw, occ_bw
def build(self):
input_queue = tf.train.string_input_producer(
[self.filenames_file], shuffle=False
)
line_reader = tf.TextLineReader()
_, line = line_reader.read(input_queue)
split_line = tf.string_split([line]).values
with tf.variable_scope("tester_dataloader_three_frames"):
with tf.variable_scope("image_reader"):
src_img_1_path = tf.string_join([self.datapath, split_line[0]])
tgt_img_path = tf.string_join([self.datapath, split_line[1]])
src_img_2_path = tf.string_join([self.datapath, split_line[2]])
src_img_1_o = self.read_image(src_img_1_path)
tgt_img_o = self.read_image(tgt_img_path)
src_img_2_o = self.read_image(src_img_2_path)
with tf.variable_scope("batch_creator"):
self.src_img_1_batch = tf.stack([src_img_1_o], 0)
self.tgt_img_batch = tf.stack([tgt_img_o], 0)
self.src_img_2_batch = tf.stack([src_img_2_o], 0)
with tf.variable_scope("shape_setter"):
self.src_img_1_batch.set_shape([1, None, None, 3])
self.tgt_img_batch.set_shape([1, None, None, 3])
self.src_img_2_batch.set_shape([1, None, None, 3])
def decode(self, serialized_example, items=None):
global _MAX_SKIP_FRAMES, _TEST_SKIP_FRAMES
features = {
'image_iter': tf.FixedLenFeature([], tf.int64),
'shape': tf.FixedLenFeature([], tf.string),
'event_count_images': tf.FixedLenFeature([], tf.string),
'event_time_images': tf.FixedLenFeature([], tf.string),
'image_times': tf.FixedLenFeature([], tf.string),
'prefix': tf.FixedLenFeature([], tf.string),
'cam': tf.FixedLenFeature([], tf.string)
}
data = tf.parse_single_example(serialized_example, features)
image_iter = data['image_iter']
prefix = data['prefix']
cam = data['cam']
image_times = tf.decode_raw(data['image_times'], tf.float64)
if self._split is 'test':
if self._skip_frames:
n_frames = _TEST_SKIP_FRAMES
else:
n_frames = 1
else:
n_frames = tf.random_uniform([],
1,
_MAX_SKIP_FRAMES,
dtype=tf.int64)
timestamps = [image_times[0], image_times[n_frames]]
event_image = self._read_events(data, n_frames)
# Get paths to grayscale png files.
prev_img_path = tf.string_join([
prefix, "/", cam, "_image",
tf.as_string(image_iter, width=5, fill='0'), ".png"
])
next_img_path = tf.string_join([
prefix, "/", cam, "_image",
tf.as_string(image_iter + n_frames * 2, width=5, fill='0'), ".png"
])
prev_image = self._read_image(prev_img_path)
next_image = self._read_image(next_img_path)
outputs = []
for item in self._items_to_features.keys():
if item == 'event_image':
outputs.append(event_image)
elif item == 'prev_image':
outputs.append(prev_image)
elif item == 'next_image':
outputs.append(next_image)
elif item == 'timestamps':
outputs.append(timestamps)
else:
raise NameError("Item {} is not valid.".format(item))
return outputs
def read_and_decode(self, filename_queue):
img1_name = tf.string_join([self.img_dir, '/', filename_queue[0]])
img2_name = tf.string_join([self.img_dir, '/', filename_queue[1]])
img1 = tf.image.decode_png(tf.read_file(img1_name), channels=3)
img1 = tf.cast(img1, tf.float32)
img2 = tf.image.decode_png(tf.read_file(img2_name), channels=3)
img2 = tf.cast(img2, tf.float32)
return img1, img2
def make_write_stage(self, ready_to_write_items):
"""
:param ready_to_write_items: a generator of (id_and_count, chunk_file_matrix, first_ordinal, num_records, record_id, namespace)
:return: a generator of (id_and_count, record_id, first_ordinal, num_records, key_basename, namespace) + (list, of, full, file, keys)
"""
write_items = sanitize_generator(
ready_to_write_items) # need to iterate multiple times
chunk_basenames = ((record_id, tf.as_string(first_ordinal))
for id_and_count, buffer_handles, first_ordinal,
num_records, record_id, namespace in write_items)
if self.overwrite:
chunk_basenames = tuple(
tf.string_join((chunk_base, ordinal_as_string),
separator="_",
name="final_chunk_join")
for chunk_base, ordinal_as_string in chunk_basenames)
else:
chunk_basenames = tuple(
tf.string_join((chunk_base, self.new_dataset_extension,
ordinal_as_string),
separator="_",
name="final_chunk_join")
for chunk_base, ordinal_as_string in chunk_basenames)
to_writer_gen = (
(key, namespace, num_records, first_ordinal, record_id,
buffer_handles)
for key, (id_and_count, buffer_handles, first_ordinal, num_records,
record_id,
namespace) in zip(chunk_basenames, write_items))
around_writer_gen = (
(id_and_count, record_id, first_ordinal, num_records, key,
namespace)
for (id_and_count, buffer_handles, first_ordinal, num_records,
record_id,
namespace), key in zip(write_items, chunk_basenames))
kwargs = {}
if self.log_goodput:
kwargs["log_directory"] = self.log_directory
kwargs["metadata"] = tuple(slice_id(a[0]) for a in write_items)
written_records = (tuple(a) for a in pipeline.ceph_write_pipeline(
upstream_tensors=to_writer_gen,
compressed=True,
user_name=self.ceph_user_name,
cluster_name=self.ceph_cluster_name,
pool_name=self.ceph_pool_name,
name="merge_ceph_write",
ceph_conf_path=str(self.ceph_conf_path),
record_types=(
get_dicts_for_extension(self.columns, text_base=False)),
**kwargs))
results = tuple(a + b
for a, b in zip(around_writer_gen, written_records))
assert len(results) == len(write_items)
return results
def make_write_stage(self, ready_to_write_items):
"""
:param ready_to_write_items: a generator of (id_and_count, chunk_file_matrix, first_ordinal, num_records, record_id, filename)
:return: a generator of (id_and_count, record_id, first_ordinal, num_records, file_basename) + (list, of, full, file, paths)
"""
write_items = tuple(
ready_to_write_items) # need to iterate multiple times
write_items = tuple(
(id_and_count, buffer_handles, first_ordinal, num_records,
record_id, file_directory)
for (id_and_count, buffer_handles, first_ordinal, num_records,
record_id), file_directory in zip((
a[:5] for a in write_items), (dirname(filename=b[5])
for b in write_items)))
chunk_basenames = (
(tf.string_join(
(file_directory, record_id),
separator=path_separator_str,
name="chunk_base_join"), tf.as_string(first_ordinal))
for id_and_count, buffer_handles, first_ordinal, num_records,
record_id, file_directory in write_items)
if self.overwrite:
chunk_basenames = tuple(
tf.string_join((chunk_base, ordinal_as_string),
separator="_",
name="final_chunk_join")
for chunk_base, ordinal_as_string in chunk_basenames)
else:
chunk_basenames = tuple(
tf.string_join((chunk_base, self.new_dataset_extension,
ordinal_as_string),
separator="_",
name="final_chunk_join")
for chunk_base, ordinal_as_string in chunk_basenames)
to_writer_gen = ((buffer_handles, record_id, first_ordinal,
num_records, file_basename)
for (id_and_count, buffer_handles, first_ordinal,
num_records, record_id), file_basename in zip((
a[:-1]
for a in write_items), chunk_basenames))
around_writer_gen = ((id_and_count, record_id, first_ordinal,
num_records, file_basename)
for (id_and_count, buffer_handles, first_ordinal,
num_records,
record_id), file_basename in zip((
a[:-1]
for a in write_items), chunk_basenames))
written_records = (
tuple(a) for a in pipeline.local_write_pipeline(
upstream_tensors=to_writer_gen,
compressed=True, # compressed controls the buffer output
record_types=(
get_dicts_for_extension(self.columns, text_base=False))))
results = tuple(a + b
for a, b in zip(around_writer_gen, written_records))
assert len(results) == len(write_items)
return results
def read_data():
img_a = tf.image.decode_image(tf.read_file(
tf.string_join(['./training_set/', self.data_queue[0]])),
channels=3)
img_b = tf.image.decode_image(tf.read_file(
tf.string_join(['./training_set/', self.data_queue[1]])),
channels=3)
img_a, img_b = preprocessing([img_a, img_b])
return img_a, img_b
def __init__(self, data_path, filenames_file, params, mode):
self.data_path = data_path
self.params = params
self.mode = mode
self.left_image_batch = None
self.right_image_batch = None
filename_queue = tf.train.string_input_producer([filenames_file],
shuffle=False)
# tf v2.1
# filename_queue = tf.data.Dataset.from_tensor_slices([filenames_file])
reader = tf.TextLineReader()
key, value = reader.read(filename_queue)
filenames = tf.string_split([value]).values
if mode == 'train':
left_image_path = tf.string_join([self.data_path, filenames[0]])
right_image_path = tf.string_join([self.data_path, filenames[1]])
left_image_o = self.read_image(left_image_path)
right_image_o = self.read_image(right_image_path)
# flip images horizontally and swap the two images with a 50% chance
if_flip = tf.random_uniform([], 0, 1) > 0.5
left_image = tf.cond(
if_flip, lambda: tf.image.flip_left_right(right_image_o),
lambda: left_image_o)
right_image = tf.cond(
if_flip, lambda: tf.image.flip_left_right(left_image_o),
lambda: right_image_o)
# do color augmentation with a 50% chance
if_augment = tf.random_uniform([], 0, 1) > 0.5
left_image, right_image = tf.cond(
if_augment,
lambda: self.do_color_augmentation(left_image, right_image),
lambda: (left_image, right_image))
left_image.set_shape([None, None, 3])
right_image.set_shape([None, None, 3])
min_after_dequeue = 2048
capacity = min_after_dequeue + 4 * params.batch_size
self.left_image_batch, self.right_image_batch = tf.train.shuffle_batch(
[left_image, right_image], params.batch_size, capacity,
min_after_dequeue, params.num_threads)
elif mode == 'test':
left_image_path = tf.string_join([self.data_path, filenames[0]])
left_image_o = self.read_image(left_image_path)
self.left_image_batch = tf.stack(
[left_image_o,
tf.image.flip_left_right(left_image_o)], 0)
self.left_image_batch.set_shape([2, None, None, 3])
def __init__(self, data_path, filenames_file, params, dataset, mode):
self.data_path = data_path
self.params = params
self.dataset = dataset
self.mode = mode
self.left_image_batch = None
self.right_image_batch = None
# 1）创建文件名队列（Queue）
input_queue = tf.train.string_input_producer([filenames_file], shuffle=False)
line_reader = tf.TextLineReader()
_, line = line_reader.read(input_queue)
split_line = tf.string_split([line]).values
# 2）编码解码
# we load only one image for test, except if we trained a stereo model
if mode == 'test' and not self.params.do_stereo:
left_image_path = tf.string_join([self.data_path, split_line[0]])
left_image_o = self.read_image(left_image_path)
else:
left_image_path = tf.string_join([self.data_path, split_line[0]])
right_image_path = tf.string_join([self.data_path, split_line[1]])
left_image_o = self.read_image(left_image_path)
right_image_o = self.read_image(right_image_path)
if mode == 'train':
# randomly flip images
# 根据随机数翻转图像（左右翻转-Width方向上），然后将右视图变成左视图
do_flip = tf.random_uniform([], 0, 1)
left_image = tf.cond(do_flip > 0.5, lambda: tf.image.flip_left_right(right_image_o), lambda: left_image_o)
right_image = tf.cond(do_flip > 0.5, lambda: tf.image.flip_left_right(left_image_o), lambda: right_image_o)
# randomly augment images
# 根据随机数进行图像增强操作
do_augment = tf.random_uniform([], 0, 1)
left_image, right_image = tf.cond(do_augment > 0.5, lambda: self.augment_image_pair(left_image, right_image), lambda: (left_image, right_image))
left_image.set_shape( [None, None, 3])
right_image.set_shape([None, None, 3])
# 3）创建样例队列
# capacity = min_after_dequeue + (num_threads + a small safety margin) * batch_size
min_after_dequeue = 2048
capacity = min_after_dequeue + 4 * params.batch_size
self.left_image_batch, self.right_image_batch = tf.train.shuffle_batch([left_image, right_image],
params.batch_size, capacity, min_after_dequeue, params.num_threads)
elif mode == 'test':
self.left_image_batch = tf.stack([left_image_o, tf.image.flip_left_right(left_image_o)], 0)
self.left_image_batch.set_shape( [2, None, None, 3])
if self.params.do_stereo:
self.right_image_batch = tf.stack([right_image_o, tf.image.flip_left_right(right_image_o)], 0)
self.right_image_batch.set_shape( [2, None, None, 3])
def __init__(self, img_dir, list1, list2, batch_size, h,w, num_threads,seed=0):
self.img_dir = img_dir
self.list1=list1
self.list2=list2
self.target_h=h
self.target_w=w
self.simg_batch = None
self.slabel_batch = None
self.timg_batch=None
self.tlabel_batch=None
queue1 = tf.train.string_input_producer([self.list1], shuffle=False)
queue2 = tf.train.string_input_producer([self.list2], shuffle=False)
# Chong: must use two seperate TextLineReader!!!
line_reader1 = tf.TextLineReader()
line_reader2 = tf.TextLineReader()
_, line1 = line_reader1.read(queue1)
_, line2 = line_reader2.read(queue2)
#pdb.set_trace()
split_line1 = tf.string_split([line1],' ').values
split_line2 = tf.string_split([line2],' ').values
p1=tf.string_join([self.img_dir,'/', split_line1[0]])
slabel=tf.string_to_number(split_line1[1],tf.int64)
p2=tf.string_join([self.img_dir,'/', split_line2[0]])
tlabel=tf.string_to_number(split_line2[1],tf.int64)
simg=self.read_image(p1)
timg=self.read_image(p2)
simg=self.augment_image(simg)
timg=self.augment_image(timg)
simg.set_shape( [10,self.target_h, self.target_w, 3])
# slabel.set_shape([None, 1])
timg.set_shape( [10,self.target_h, self.target_w, 3])
# tlabel.set_shape([None,1])
simg=self.pre_process(simg)
timg=self.pre_process(timg)
simg=simg[...,::-1]
timg=timg[...,::-1]
# capacity = min_after_dequeue + (num_threads + a small safety margin) * batch_size
min_after_dequeue = 2048
capacity = min_after_dequeue + (4+num_threads)* batch_size
simg_batch, self.slabel_batch, timg_batch, self.tlabel_batch = tf.train.shuffle_batch([simg,slabel,timg,tlabel],
batch_size, capacity, min_after_dequeue, num_threads)
self.simg_batch=tf.squeeze(simg_batch,[0])
self.timg_batch=tf.squeeze(timg_batch,[0])
def read_data():
img_a = tf.image.decode_image(tf.read_file(tf.string_join(['/home/opt603/high_io_data/GOPRO_Large/train/', self.data_queue[0]])),
channels=3)
img_b = tf.image.decode_image(tf.read_file(tf.string_join(['/home/opt603/high_io_data/GOPRO_Large/train/', self.data_queue[1]])),
channels=3)
img_c = tf.image.decode_image(tf.read_file(tf.string_join(['/home/opt603/high_io_data/GOPRO_Large/train/', self.data_queue[2]])),
channels=3)
img_a, img_b, img_c = preprocessing([img_a, img_b, img_c])
img_c = tf.image.rgb_to_grayscale(img_c)
return img_a, img_b, img_c
def parse_function_train(self, line):
split_line = tf.string_split([line]).values
image_path = tf.string_join([self.data_path, split_line[0]])
depth_gt_path = tf.string_join(
[self.gt_path, tf.string_strip(split_line[1])])
if self.params.dataset == 'nyu':
image = tf.image.decode_jpeg(tf.read_file(image_path))
else:
image = tf.image.decode_png(tf.read_file(image_path))
depth_gt = tf.image.decode_png(tf.read_file(depth_gt_path),
channels=0,
dtype=tf.uint16)
if self.params.dataset == 'nyu':
depth_gt = tf.cast(depth_gt, tf.float32) / 1000.0
else:
depth_gt = tf.cast(depth_gt, tf.float32) / 256.0
image = tf.image.convert_image_dtype(image, tf.float32)
focal = tf.string_to_number(split_line[2])
# To avoid blank boundaries due to pixel registration
if self.params.dataset == 'nyu':
depth_gt = depth_gt[45:472, 43:608, :]
image = image[45:472, 43:608, :]
if self.do_kb_crop is True:
print('Cropping training images as kitti benchmark images')
height = tf.shape(image)[0]
width = tf.shape(image)[1]
top_margin = tf.to_int32(height - 352)
left_margin = tf.to_int32((width - 1216) / 2)
depth_gt = depth_gt[top_margin:top_margin + 352,
left_margin:left_margin + 1216, :]
image = image[top_margin:top_margin + 352,
left_margin:left_margin + 1216, :]
if self.do_rotate is True:
random_angle = tf.random_uniform([], -self.degree * 3.141592 / 180,
self.degree * 3.141592 / 180)
image = tf.contrib.image.rotate(image,
random_angle,
interpolation='BILINEAR')
depth_gt = tf.contrib.image.rotate(depth_gt,
random_angle,
interpolation='NEAREST')
print('Do random cropping from fixed size input')
image, depth_gt = self.random_crop_fixed_size(image, depth_gt)
return image, depth_gt, focal
def read_data():
img_a = tf.image.decode_image(tf.read_file(
tf.string_join([
'/home/opt603/sda5/GOPRO_Large/train/', self.data_queue[0]
])),
channels=3)
img_b = tf.image.decode_image(tf.read_file(
tf.string_join([
'/home/opt603/sda5/GOPRO_Large/train/', self.data_queue[1]
])),
channels=3)
img_a, img_b = preprocessing([img_a, img_b])
return img_a, img_b
def read_fix_size_image_format(dataset_root_dir_string, filename_queue):
# http://stackoverflow.com/questions/37198357/loading-images-and-labels-from-csv-file-using-tensorflow
# https://www.tensorflow.org/versions/r0.11/how_tos/reading_data/index.html#file-formats
reader = tf.TextLineReader() # 每次读取一行csv文件
key, value = reader.read(filename_queue)
# Default values
record_defaults = [[''], ['']]
# 使用tf.decode_csv来对每一行进行解析，分割为图像路径和标签路径
image_path, annotation_path = tf.decode_csv(
value, record_defaults=record_defaults)
# csv 里保存的不是绝对路径，需要和 dataset_root_dir_string 一起拼装成完整的路径
image_path = tf.string_join(
[tf.constant(dataset_root_dir_string), image_path])
annotation_path = tf.string_join(
[tf.constant(dataset_root_dir_string), annotation_path])
# 对最后的图像和标签路径进行读取
image_content = tf.read_file(image_path)
annotation_content = tf.read_file(annotation_path)
# http://stackoverflow.com/questions/34746777/why-do-i-get-valueerror-image-must-be-fully-defined-when-transforming-im
# http://stackoverflow.com/questions/37772329/tensorflow-tensor-set-shape-valueerror-image-must-be-fully-defined
# image is jpg, annotation is png
# 对jpg格式和png格式的图像进行解码，得到图像的像素值，
# 这个像素值可以用于显示图像。如果没有解码，读取的图像是一个字符串
image_tensor = tf.image.decode_jpeg(image_content, channels=3)
annotation_tensor = tf.image.decode_png(annotation_content, channels=1)
# decode之后，一定要设置 image 的大小，或者 resize 到一个size，否则会 crash
image_tensor = tf.reshape(image_tensor,
[const.image_height, const.image_width, 3])
annotation_tensor = tf.reshape(annotation_tensor,
[const.image_height, const.image_width, 1])
image_float = tf.to_float(image_tensor)
annotation_float = tf.to_float(annotation_tensor)
# print('debug, image_float shape is: {}'.format(image_float.get_shape()))
# print('debug, annotation_float shape is: {}'.format(annotation_float.get_shape()))
if const.use_batch_norm == True:
image_float = image_float / 255.0 # 进行归一化处理
else:
# 这个分支主要是为了匹配不使用 batch norm 时的 VGG
image_float = mean_image_subtraction(
image_float, [R_MEAN, G_MEAN, B_MEAN]) # 一个不做归一化，一个做归一化处理
# 不管是不是 VGG，annotation 都需要归一化
annotation_float = annotation_float / 255.0
return image_float, annotation_float, image_path
def add_inst_stats(correct1, correct5, instcombs, train_test_table_selector):
with tf.name_scope('instStats') as scope:
sorted_labels = tf.py_func(_sort_labels, [instcombs], tf.string)
update_train_top1_table, export_train_top1_table = _add_onehot_to_hash_table_average(
correct1, sorted_labels)
update_train_top5_table, export_train_top5_table = _add_onehot_to_hash_table_average(
correct5, sorted_labels)
update_test_top1_table, export_test_top1_table = _add_onehot_to_hash_table_average(
correct1, sorted_labels)
update_test_top5_table, export_test_top5_table = _add_onehot_to_hash_table_average(
correct5, sorted_labels)
update_top1_table, export_top1_table = tf.cond(
tf.equal(train_test_table_selector, 0),
lambda: [update_train_top1_table, export_train_top1_table],
lambda: [update_test_top1_table, export_test_top1_table])
update_top5_table, export_top5_table = tf.cond(
tf.equal(train_test_table_selector, 0),
lambda: [update_train_top5_table, export_train_top5_table],
lambda: [update_test_top5_table, export_test_top5_table])
# top1_plot_op = tfplot.plot(_create_bar_stats_figure, [export_top1_table[0], export_top1_table[1]])
# top5_plot_op = tfplot.plot(_create_bar_stats_figure, [export_top5_table[0], export_top5_table[1]])
# tf.summary.image('top1_instcomb', tf.expand_dims(top1_plot_op, 0), max_outputs=1)
# tf.summary.image('top5_instcomb', tf.expand_dims(top5_plot_op, 0), max_outputs=1)
top_sorted = tf.nn.top_k(export_top1_table[1],
k=tf.shape(export_top1_table[1])[0],
sorted=True)
sidx = top_sorted.indices
tf.summary.text(
'top1_instcomb',
tf.string_join([
tf.gather(export_top1_table[0], sidx),
tf.as_string(tf.gather(export_top1_table[1], sidx))
],
separator=' - '))
tf.summary.text(
'top5_instcomb',
tf.string_join([
tf.gather(export_top5_table[0], sidx),
tf.as_string(tf.gather(export_top5_table[1], sidx))
],
separator=' - '))
reset_tables = tf.local_variables_initializer()
return [update_top1_table, update_top5_table], reset_tables
def _parse_image(folder_t, frame_t, size_t):
image_filename = tf.string_join([
dataset_data_dir, folder_t,
tf.string_join(
[tf.as_string(frame_t[0], width=6, fill='0'), '.JPEG'])
],
separator='/')
image_raw = tf.read_file(image_filename)
image_decoded = tf.image.decode_jpeg(image_raw, channels=3)
image_decoded = tf.image.convert_image_dtype(image_decoded,
dtype=tf.float32)
image_resized = tf.image.resize_images(image_decoded, [360, 480])
# image_resized = image_decoded
return image_resized
def parse_example(parsed_features):
label = parsed_features['comb/label']
ins = parsed_features['example/input']
inst1 = tf.regex_replace(parsed_features['comb/inst1'], ' ', '_')
inst2 = tf.regex_replace(parsed_features['comb/inst2'], ' ', '_')
type1 = parsed_features['comb/type1']
type2 = parsed_features['comb/type2']
file1 = parsed_features['comb/file1']
file2 = parsed_features['comb/file2']
genre = parsed_features['comb/genre']
id = parsed_features['comb/id']
return ins, label, tf.string_join([type1, ' x ', type2]), tf.string_join(
[inst1, ' x ',
inst2]), genre, id, tf.string_join([file1, ' x ', file2])
def _read_from_disk_temporal(
fpath, nframes, num_samples=25,
optical_flow_frames=10, start_frame=0,
file_prefix='', file_zero_padding=4, file_index=1,
dataset_dir='', step=None):
duration = nframes
if step is None:
if num_samples == 1:
step = tf.random_uniform([1], 0, nframes-optical_flow_frames-1, dtype='int32')[0]
else:
step = tf.cast((duration-tf.constant(optical_flow_frames)) /
(tf.constant(num_samples)), 'int32')
allimgs = []
with tf.variable_scope('read_flow_video'):
for i in range(num_samples):
if num_samples == 1:
i = 1 # so that the random step value can be used
with tf.variable_scope('read_flow_image'):
flow_img = []
for j in range(optical_flow_frames):
with tf.variable_scope('read_flow_channels'):
for dr in ['x', 'y']:
prefix = file_prefix + '_' if file_prefix else ''
impath = tf.string_join([
tf.constant(dataset_dir + '/'),
fpath, tf.constant('/'),
prefix, '%s_' % dr,
tf.as_string(start_frame + i * step + file_index + j,
width=file_zero_padding, fill='0'),
tf.constant('.jpg')])
img_str = tf.read_file(impath)
flow_img.append(img_str)
allimgs.append(flow_img)
return allimgs
def simple_example(step):
# Text summaries log arbitrary text. This can be encoded with ASCII or
# UTF-8. Here's a simple example, wherein we greet the user on each
# step:
step_string = tf.as_string(step)
greeting = tf.string_join(['Hello from step ', step_string, '!'])
tf.summary.text('greeting', greeting)
def _read_from_disk_spatial(fpath, nframes, num_samples=25, start_frame=0,
file_prefix='', file_zero_padding=4, file_index=1,
dataset_dir='', step=None):
duration = nframes
if step is None:
if num_samples == 1:
step = tf.random_uniform([1], 0, nframes, dtype='int32')[0]
else:
step = tf.cast((duration-tf.constant(1)) /
(tf.constant(num_samples-1)), 'int32')
allimgs = []
with tf.variable_scope('read_rgb_video'):
for i in range(num_samples):
if num_samples == 1:
i = 1 # so that the random step value can be used
with tf.variable_scope('read_rgb_image'):
prefix = file_prefix + '_' if file_prefix else ''
impath = tf.string_join([
tf.constant(dataset_dir + '/'),
fpath, tf.constant('/'),
prefix,
tf.as_string(start_frame + i * step + file_index,
width=file_zero_padding, fill='0'),
tf.constant('.jpg')])
img_str = tf.read_file(impath)
allimgs.append(img_str)
return allimgs
def higher_order_tensors(step):
# We're not limited to passing scalar tensors to the summary
# operation. If we pass a rank-1 or rank-2 tensor, it'll be visualized
# as a table in TensorBoard. (For higher-ranked tensors, you'll see
# just a 2D slice of the data.)
#
# To demonstrate this, let's create a multiplication table.
# First, we'll create the table body, a `step`-by-`step` array of
# strings.
numbers = tf.range(step)
numbers_row = tf.expand_dims(numbers, 0) # shape: [1, step]
numbers_column = tf.expand_dims(numbers, 1) # shape: [step, 1]
products = tf.matmul(numbers_column, numbers_row) # shape: [step, step]
table_body = tf.as_string(products)
# Next, we'll create a header row and column, and a little
# multiplication sign to put in the corner.
bold_numbers = tf.string_join(['**', tf.as_string(numbers), '**'])
bold_row = tf.expand_dims(bold_numbers, 0)
bold_column = tf.expand_dims(bold_numbers, 1)
corner_cell = tf.constant(u'\u00d7'.encode('utf-8')) # MULTIPLICATION SIGN
# Now, we have to put the pieces together. Using `axis=0` stacks
# vertically; using `axis=1` juxtaposes horizontally.
table_body_and_top_row = tf.concat([bold_row, table_body], axis=0)
table_left_column = tf.concat([[[corner_cell]], bold_column], axis=0)
table_full = tf.concat([table_left_column, table_body_and_top_row], axis=1)
# The result, `table_full`, is a rank-2 string tensor of shape
# `[step + 1, step + 1]`. We can pass it directly to the summary, and
# we'll get a nicely formatted table in TensorBoard.
tf.summary.text('multiplication_table', table_full)
def testStateSaverScopeNames(self):
batch_size = tf.constant(2)
sqss_scope_name = "unique_scope_name_for_sqss"
num_unroll = 2
length = 3
key = tf.string_join(["key_", tf.as_string(tf.cast(
10000 * tf.random_uniform(()), tf.int32))])
padded_length = 4
sequences = {"seq1": np.random.rand(padded_length, 5),
"seq2": np.random.rand(padded_length, 4, 2)}
context = {"context1": [3, 4]}
initial_states = {"state1": np.random.rand(6, 7),
"state2": np.random.rand(8)}
state_saver = tf.contrib.training.SequenceQueueingStateSaver(
batch_size=batch_size,
num_unroll=num_unroll,
input_length=length,
input_key=key,
input_sequences=sequences,
input_context=context,
initial_states=initial_states,
name=sqss_scope_name)
prefetch_op = state_saver.prefetch_op
next_batch = state_saver.next_batch
self.assertTrue(state_saver.barrier.barrier_ref.name.startswith(
"%s/" % sqss_scope_name))
self.assertTrue(prefetch_op.name.startswith("%s/" % sqss_scope_name))
self.assertTrue(next_batch.key.name.startswith("%s/" % sqss_scope_name))
def add_distance_transform(tensors, labels, distance_transform_fn):
args_list = [
tensors["unnormalized_img"], tensors["label"], tensors["raw_label"],
labels[Constants.STRATEGY], labels[Constants.IGNORE_CLASSES]
]
if "old_label" in tensors:
args_list.append(tensors["old_label"])
u0, u1, num_clicks = tf.py_func(distance_transform_fn,
args_list,
[tf.float32, tf.float32, tf.int64],
name="create_distance_transform")
u0 = tf.expand_dims(u0, axis=2)
u0.set_shape(tensors["unnormalized_img"].get_shape().as_list()[:-1] + [1])
u1 = tf.expand_dims(u1, axis=2)
u1.set_shape(tensors["unnormalized_img"].get_shape().as_list()[:-1] + [1])
shape = tensors["tag"].get_shape()
im_path = tf.string_join(
[tensors["tag"], tf.as_string(num_clicks)],
separator=":",
name="JoinPath")
im_path.set_shape(shape)
tensors[Constants.DT_NEG] = u0
tensors[Constants.DT_POS] = u1
tensors["tag"] = im_path
return tensors
def get_test_batches(images, batch_size=32, num_epochs=None):
"""Returns a batch of images from test set."""
height = 512
width = 512
num_channels = 3
filename = tf.train.slice_input_producer([images],
shuffle=False,
num_epochs=num_epochs)[0]
prefix = os.path.join(os.getcwd(), "data", "test", "")
filename = tf.string_join([prefix, filename])
image = tf.read_file(filename)
image = tf.image.decode_jpeg(image, channels=num_channels)
image = tf.image.resize_images(image, [height, width])
img_batch = tf.train.batch(
[image],
batch_size,
capacity=4000,
allow_smaller_final_batch=True,
num_threads=4,
)
return img_batch
def generate_run(self, run_name, include_graph):
"""Create a run with a text summary, metadata, and optionally a graph."""
tf.reset_default_graph()
k1 = tf.constant(math.pi, name='k1')
k2 = tf.constant(math.e, name='k2')
result = (k1 ** k2) - k1
expected = tf.constant(20.0, name='expected')
error = tf.abs(result - expected, name='error')
message_prefix_value = 'error ' * 1000
true_length = len(message_prefix_value)
assert true_length > self._MESSAGE_PREFIX_LENGTH_LOWER_BOUND, true_length
message_prefix = tf.constant(message_prefix_value, name='message_prefix')
error_message = tf.string_join([message_prefix,
tf.as_string(error, name='error_string')],
name='error_message')
summary_message = tf.summary.text('summary_message', error_message)
sess = tf.Session()
writer = tf.summary.FileWriter(os.path.join(self.logdir, run_name))
if include_graph:
writer.add_graph(sess.graph)
options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
s = sess.run(summary_message, options=options, run_metadata=run_metadata)
writer.add_summary(s)
writer.add_run_metadata(run_metadata, self._METADATA_TAG)
writer.close()
def __init__(self, data_path, filenames_file, params, dataset, mode):
self.data_path = data_path
self.params = params
self.dataset = dataset
self.mode = mode
self.left_image_batch = None
self.right_image_batch = None
input_queue = tf.train.string_input_producer([filenames_file], shuffle=False)
line_reader = tf.TextLineReader()
_, line = line_reader.read(input_queue)
split_line = tf.string_split([line]).values
# we load only one image for test, except if we trained a stereo model
if mode == 'test' and not self.params.do_stereo:
left_image_path = tf.string_join([self.data_path, split_line[0]])
left_image_o = self.read_image(left_image_path)
if mode == 'test':
self.left_image_batch = tf.stack([left_image_o, tf.image.flip_left_right(left_image_o)], 0)
self.left_image_batch.set_shape( [2, None, None, 3])
if self.params.do_stereo:
self.right_image_batch = tf.stack([right_image_o, tf.image.flip_left_right(right_image_o)], 0)
self.right_image_batch.set_shape( [2, None, None, 3])
def main(data_path, batch_size, latent_dim, steps, examples_limit, log_dir,
gradient_penalty, generator_lr, discriminator_lr, checkpoint_freq):
"""Trains a DCGAN to generate CNN training images."""
# Load data
real = load_image_h5(data_path)
# Process real data
real_data = np.expand_dims(real[:, :, :, 0], 3)
real_data = (real_data.astype('float32') - 127.5) / 127.5
real_data = real_data[:examples_limit]
# Create data iterator
data_iterator, iterator_init_hook = data_to_stream(real_data, batch_size)
# Configure GAN model
gan_model = tfgan.gan_model(generator_fn=networks.generator,
discriminator_fn=networks.discriminator,
real_data=data_iterator.get_next(),
generator_inputs=tf.random_normal(
[batch_size, latent_dim]))
tfgan.eval.add_gan_model_image_summaries(gan_model)
tfgan.eval.add_regularization_loss_summaries(gan_model)
# Set up loss functions
gan_loss = tfgan.gan_loss(gan_model,
gradient_penalty_weight=gradient_penalty,
add_summaries=True)
# Configure training ops
generator_opt = tf.train.AdamOptimizer(generator_lr, beta1=0.5)
discriminator_opt = tf.train.AdamOptimizer(discriminator_lr, beta1=0.5)
train_ops = tfgan.gan_train_ops(gan_model,
gan_loss,
generator_optimizer=generator_opt,
discriminator_optimizer=discriminator_opt,
summarize_gradients=True)
status_message = tf.string_join([
'Starting train step: ',
tf.as_string(tf.train.get_or_create_global_step())
],
name='status_message')
if os.path.exists(log_dir):
print(
'Log directory already exists. Exiting to avoid overwriting other model.'
)
exit(0)
# Begin training
tfgan.gan_train(train_ops,
logdir=log_dir,
save_checkpoint_secs=checkpoint_freq,
hooks=[
tf.train.StopAtStepHook(num_steps=steps),
tf.train.LoggingTensorHook([status_message],
every_n_iter=100),
iterator_init_hook
])
def make_status_message(model):
"""Makes a string `Tensor` of training status."""
return tf.string_join(
[
'Starting train step: current_image_id: ',
tf.as_string(model.current_image_id), ', progress: ',
tf.as_string(model.progress), ', num_blocks: {}'.format(
model.num_blocks), ', batch_size: {}'.format(model.batch_size)
],
name='status_message')
def setUp(self):
super(BatchSequencesWithStatesTest, self).setUp()
self.value_length = 4
self.batch_size = 2
self.key = tf.string_join(["key_", tf.as_string(tf.cast(
10000 * tf.random_uniform(()), tf.int32))])
self.sequences = {"seq1": np.random.rand(self.value_length, 5),
"seq2": np.random.rand(self.value_length, 4, 2)}
self.context = {"context1": [3, 4]}
self.initial_states = {"state1": np.random.rand(6, 7),
"state2": np.random.rand(8)}
def main(_):
# Create the log_dir if not exist.
if not tf.gfile.Exists(FLAGS.train_log_dir):
tf.gfile.MakeDirs(FLAGS.train_log_dir)
# Shard the model to different parameter servers.
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Create the input dataset.
with tf.name_scope('inputs'):
images, labels = data_provider.provide_data(
FLAGS.image_file_patterns, FLAGS.batch_size, FLAGS.patch_size)
# Define the model.
with tf.name_scope('model'):
model = _define_model(images, labels)
# Add image summary.
tfgan.eval.add_stargan_image_summaries(
model,
num_images=len(FLAGS.image_file_patterns) * FLAGS.batch_size,
display_diffs=True)
# Define the model loss.
loss = tfgan.stargan_loss(model)
# Define the train ops.
with tf.name_scope('train_ops'):
train_ops = _define_train_ops(model, loss)
# Define the train steps.
train_steps = _define_train_step()
# Define a status message.
status_message = tf.string_join(
[
'Starting train step: ',
tf.as_string(tf.train.get_or_create_global_step())
],
name='status_message')
# Train the model.
tfgan.gan_train(
train_ops,
FLAGS.train_log_dir,
get_hooks_fn=tfgan.get_sequential_train_hooks(train_steps),
hooks=[
tf.train.StopAtStepHook(num_steps=FLAGS.max_number_of_steps),
tf.train.LoggingTensorHook([status_message], every_n_iter=10)
],
master=FLAGS.master,
is_chief=FLAGS.task == 0)
def markdown_table(step):
# The text summary can also contain Markdown, including Markdown
# tables. Markdown tables look like this:
#
# | hello | there |
# |-------|-------|
# | this | is |
# | a | table |
#
# The leading and trailing pipes in each row are optional, and the text
# doesn't actually have to be neatly aligned, so we can create these
# pretty easily. Let's do so.
header_row = 'Pounds of chocolate | Happiness'
chocolate = tf.range(step)
happiness = tf.square(chocolate + 1)
chocolate_column = tf.as_string(chocolate)
happiness_column = tf.as_string(happiness)
table_rows = tf.string_join([chocolate_column, " | ", happiness_column])
table_body = tf.reduce_join(table_rows, separator='\n')
table = tf.string_join([header_row, "---|---", table_body], separator='\n')
preamble = 'We conducted an experiment and found the following data:\n\n'
result = tf.string_join([preamble, table])
tf.summary.text('chocolate_study', result)
def main(_):
if not tf.gfile.Exists(FLAGS.train_log_dir):
tf.gfile.MakeDirs(FLAGS.train_log_dir)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
with tf.name_scope('inputs'):
images_x, images_y = data_provider.provide_custom_data(
[FLAGS.image_set_x_file_pattern, FLAGS.image_set_y_file_pattern],
batch_size=FLAGS.batch_size,
patch_size=FLAGS.patch_size)
# Set batch size for summaries.
images_x.set_shape([FLAGS.batch_size, None, None, None])
images_y.set_shape([FLAGS.batch_size, None, None, None])
# Define CycleGAN model.
cyclegan_model = _define_model(images_x, images_y)
# Define CycleGAN loss.
cyclegan_loss = tfgan.cyclegan_loss(
cyclegan_model,
cycle_consistency_loss_weight=FLAGS.cycle_consistency_loss_weight,
tensor_pool_fn=tfgan.features.tensor_pool)
# Define CycleGAN train ops.
train_ops = _define_train_ops(cyclegan_model, cyclegan_loss)
# Training
train_steps = tfgan.GANTrainSteps(1, 1)
status_message = tf.string_join(
[
'Starting train step: ',
tf.as_string(tf.train.get_or_create_global_step())
],
name='status_message')
if not FLAGS.max_number_of_steps:
return
tfgan.gan_train(
train_ops,
FLAGS.train_log_dir,
get_hooks_fn=tfgan.get_sequential_train_hooks(train_steps),
hooks=[
tf.train.StopAtStepHook(num_steps=FLAGS.max_number_of_steps),
tf.train.LoggingTensorHook([status_message], every_n_iter=10)
],
master=FLAGS.master,
is_chief=FLAGS.task == 0)
def clip_to_waveform(clip, clip_dir=None):
"""Decodes a WAV clip into a waveform tensor."""
'''
data , sampling_rate = librosa.load('data/sound.wav', sr=SAMPLE_RATE)
# for use in tensorflow
data_tensor = tf.convert_to_tensor( data )
'''
# Decode the WAV-format clip into a waveform tensor where
# the values lie in [-1, +1].
clip_path = tf.string_join([clip_dir, clip], separator=os.sep)
clip_data = tf.read_file(clip_path)
waveform, sr = tf_audio.decode_wav(clip_data)
# Assert that the clip has the expected sample rate.
check_sr = tf.assert_equal(sr, sr)
# and check that it is mono.
check_channels = tf.assert_equal(tf.shape(waveform)[1], 1)
with tf.control_dependencies([tf.group(check_sr, check_channels)]):
return tf.squeeze(waveform)
def __init__(self, config, batch_size, one_hot=False):
self.lookup = None
reader = tf.TextLineReader()
filename_queue = tf.train.string_input_producer(["chargan.txt"])
key, x = reader.read(filename_queue)
vocabulary = self.get_vocabulary()
table = tf.contrib.lookup.string_to_index_table_from_tensor(
mapping = vocabulary, default_value = 0)
x = tf.string_join([x, tf.constant(" " * 64)])
x = tf.substr(x, [0], [64])
x = tf.string_split(x,delimiter='')
x = tf.sparse_tensor_to_dense(x, default_value=' ')
x = tf.reshape(x, [64])
x = table.lookup(x)
self.one_hot = one_hot
if one_hot:
x = tf.one_hot(x, len(vocabulary))
x = tf.cast(x, dtype=tf.float32)
x = tf.reshape(x, [1, int(x.get_shape()[0]), int(x.get_shape()[1]), 1])
else:
x = tf.cast(x, dtype=tf.float32)
x -= len(vocabulary)/2.0
x /= len(vocabulary)/2.0
x = tf.reshape(x, [1,1, 64, 1])
num_preprocess_threads = 8
x = tf.train.shuffle_batch(
[x],
batch_size=batch_size,
num_threads=num_preprocess_threads,
capacity= 5000,
min_after_dequeue=500,
enqueue_many=True)
self.x = x
self.table = table
def testStringJoin(self):
input0 = ["a", "b"]
input1 = "a"
input2 = [["b"], ["c"]]
with self.test_session():
output = tf.string_join([input0, input1])
self.assertAllEqual(output.eval(), [b"aa", b"ba"])
output = tf.string_join([input0, input1], separator="--")
self.assertAllEqual(output.eval(), [b"a--a", b"b--a"])
output = tf.string_join([input0, input1, input0], separator="--")
self.assertAllEqual(output.eval(), [b"a--a--a", b"b--a--b"])
output = tf.string_join([input1] * 4, separator="!")
self.assertEqual(output.eval(), b"a!a!a!a")
output = tf.string_join([input2] * 2, separator="")
self.assertAllEqual(output.eval(), [[b"bb"], [b"cc"]])
with self.assertRaises(ValueError): # Inconsistent shapes
tf.string_join([input0, input2]).eval()
def run(logdir, run_name, wave_name, wave_constructor):
"""Generate wave data of the given form.
The provided function `wave_constructor` should accept a scalar tensor
of type float32, representing the frequency (in Hz) at which to
construct a wave, and return a tensor of shape [1, _samples(), `n`]
representing audio data (for some number of channels `n`).
Waves will be generated at frequencies ranging from A4 to A5.
Arguments:
logdir: the top-level directory into which to write summary data
run_name: the name of this run; will be created as a subdirectory
under logdir
wave_name: the name of the wave being generated
wave_constructor: see above
"""
tf.reset_default_graph()
tf.set_random_seed(0)
# On each step `i`, we'll set this placeholder to `i`. This allows us
# to know "what time it is" at each step.
step_placeholder = tf.placeholder(tf.float32, shape=[])
# We want to linearly interpolate a frequency between A4 (440 Hz) and
# A5 (880 Hz).
with tf.name_scope('compute_frequency'):
f_min = 440.0
f_max = 880.0
t = step_placeholder / (FLAGS.steps - 1)
frequency = f_min * (1.0 - t) + f_max * t
# Let's log this frequency, just so that we can make sure that it's as
# expected.
tf.summary.scalar('frequency', frequency)
# Now, we pass this to the wave constructor to get our waveform. Doing
# so within a name scope means that any summaries that the wave
# constructor produces will be namespaced.
with tf.name_scope(wave_name):
waveform = wave_constructor(frequency)
# We also have the opportunity to annotate each audio clip with a
# label. This is a good place to include the frequency, because it'll
# be visible immediately next to the audio clip.
with tf.name_scope('compute_labels'):
samples = tf.shape(waveform)[0]
wave_types = tf.tile(["*Wave type:* `%s`." % wave_name], [samples])
frequencies = tf.string_join([
"*Frequency:* ",
tf.tile([tf.as_string(frequency, precision=2)], [samples]),
" Hz.",
])
samples = tf.string_join([
"*Sample:* ", tf.as_string(tf.range(samples) + 1),
" of ", tf.as_string(samples), ".",
])
labels = tf.string_join([wave_types, frequencies, samples], separator=" ")
# We can place a description next to the summary in TensorBoard. This
# is a good place to explain what the summary represents, methodology
# for creating it, etc. Let's include the source code of the function
# that generated the wave.
source = '\n'.join(' %s' % line.rstrip()
for line in inspect.getsourcelines(wave_constructor)[0])
description = ("A wave of type `%r`, generated via:\n\n%s"
% (wave_name, source))
# Here's the crucial piece: we interpret this result as audio.
summary.op('waveform', waveform, FLAGS.sample_rate,
labels=labels,
display_name=wave_name,
description=description)
# Now, we can collect up all the summaries and begin the run.
summ = tf.summary.merge_all()
sess = tf.Session()
writer = tf.summary.FileWriter(os.path.join(logdir, run_name))
writer.add_graph(sess.graph)
sess.run(tf.global_variables_initializer())
for step in xrange(FLAGS.steps):
s = sess.run(summ, feed_dict={step_placeholder: float(step)})
writer.add_summary(s, global_step=step)
writer.close()
def main(_):
if not tf.gfile.Exists(FLAGS.train_log_dir):
tf.gfile.MakeDirs(FLAGS.train_log_dir)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Get real and distorted images.
with tf.device('/cpu:0'), tf.name_scope('inputs'):
real_images = data_provider.provide_data(
'train', FLAGS.batch_size, dataset_dir=FLAGS.dataset_dir,
patch_size=FLAGS.patch_size)
distorted_images = _distort_images(
real_images, downscale_size=int(FLAGS.patch_size / 2),
upscale_size=FLAGS.patch_size)
# Create a GANModel tuple.
gan_model = tfgan.gan_model(
generator_fn=networks.generator,
discriminator_fn=networks.discriminator,
real_data=real_images,
generator_inputs=distorted_images)
tfgan.eval.add_image_comparison_summaries(
gan_model, num_comparisons=3, display_diffs=True)
tfgan.eval.add_gan_model_image_summaries(gan_model, grid_size=3)
# Define the GANLoss tuple using standard library functions.
with tf.name_scope('losses'):
gan_loss = tfgan.gan_loss(
gan_model,
generator_loss_fn=tfgan.losses.least_squares_generator_loss,
discriminator_loss_fn=tfgan.losses.least_squares_discriminator_loss)
# Define the standard L1 pixel loss.
l1_pixel_loss = tf.norm(gan_model.real_data - gan_model.generated_data,
ord=1) / FLAGS.patch_size ** 2
# Modify the loss tuple to include the pixel loss. Add summaries as well.
gan_loss = tfgan.losses.combine_adversarial_loss(
gan_loss, gan_model, l1_pixel_loss,
weight_factor=FLAGS.weight_factor)
with tf.name_scope('train_ops'):
# Get the GANTrain ops using the custom optimizers and optional
# discriminator weight clipping.
gen_lr, dis_lr = _lr(FLAGS.generator_lr, FLAGS.discriminator_lr)
gen_opt, dis_opt = _optimizer(gen_lr, dis_lr)
train_ops = tfgan.gan_train_ops(
gan_model,
gan_loss,
generator_optimizer=gen_opt,
discriminator_optimizer=dis_opt,
summarize_gradients=True,
colocate_gradients_with_ops=True,
aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N,
transform_grads_fn=tf.contrib.training.clip_gradient_norms_fn(1e3))
tf.summary.scalar('generator_lr', gen_lr)
tf.summary.scalar('discriminator_lr', dis_lr)
# Use GAN train step function if using adversarial loss, otherwise
# only train the generator.
train_steps = tfgan.GANTrainSteps(
generator_train_steps=1,
discriminator_train_steps=int(FLAGS.weight_factor > 0))
# Run the alternating training loop. Skip it if no steps should be taken
# (used for graph construction tests).
status_message = tf.string_join(
['Starting train step: ',
tf.as_string(tf.train.get_or_create_global_step())],
name='status_message')
if FLAGS.max_number_of_steps == 0: return
tfgan.gan_train(
train_ops,
FLAGS.train_log_dir,
get_hooks_fn=tfgan.get_sequential_train_hooks(train_steps),
hooks=[tf.train.StopAtStepHook(num_steps=FLAGS.max_number_of_steps),
tf.train.LoggingTensorHook([status_message], every_n_iter=10)],
master=FLAGS.master,
is_chief=FLAGS.task == 0)
def testStateSaverWithTwoSimpleSteps(self):
with self.test_session() as sess:
batch_size_value = 2
batch_size = tf.constant(batch_size_value)
num_unroll = 2
length = 3
key = tf.string_join(["key_", tf.as_string(tf.cast(
10000 * tf.random_uniform(()), tf.int32))])
padded_length = 4
sequences = {"seq1": np.random.rand(padded_length, 5),
"seq2": np.random.rand(padded_length, 4, 2)}
context = {"context1": [3, 4]}
initial_states = {"state1": np.random.rand(6, 7),
"state2": np.random.rand(8)}
state_saver = tf.contrib.training.SequenceQueueingStateSaver(
batch_size=batch_size,
num_unroll=num_unroll,
input_length=length,
input_key=key,
input_sequences=sequences,
input_context=context,
initial_states=initial_states)
initial_key_value_0, _ = sess.run((key, state_saver.prefetch_op))
initial_key_value_1, _ = sess.run((key, state_saver.prefetch_op))
initial_key_value_0 = initial_key_value_0.decode("ascii")
initial_key_value_1 = initial_key_value_1.decode("ascii")
# Step 1
next_batch = state_saver.next_batch
(key_value, next_key_value, seq1_value, seq2_value, context1_value,
state1_value, state2_value, length_value, _, _) = sess.run(
(next_batch.key,
next_batch.next_key,
next_batch.sequences["seq1"],
next_batch.sequences["seq2"],
next_batch.context["context1"],
next_batch.state("state1"),
next_batch.state("state2"),
next_batch.length,
next_batch.save_state("state1", next_batch.state("state1") + 1),
next_batch.save_state("state2", next_batch.state("state2") - 1)))
expected_first_keys = set(
("00000_of_00002:%s" % x).encode("ascii")
for x in (initial_key_value_0, initial_key_value_1))
expected_second_keys = set(
("00001_of_00002:%s" % x).encode("ascii")
for x in (initial_key_value_0, initial_key_value_1))
expected_final_keys = set(
("STOP:%s" % x).encode("ascii")
for x in (initial_key_value_0, initial_key_value_1))
self.assertEqual(set(key_value), expected_first_keys)
self.assertEqual(set(next_key_value), expected_second_keys)
self.assertAllEqual(
context1_value, np.tile(context["context1"], (batch_size_value, 1)))
self.assertAllEqual(
seq1_value, np.tile(sequences["seq1"][np.newaxis, 0:2, :],
(batch_size_value, 1, 1)))
self.assertAllEqual(
seq2_value,
np.tile(sequences["seq2"][np.newaxis, 0:2, :, :],
(batch_size_value, 1, 1, 1)))
self.assertAllEqual(
state1_value,
np.tile(initial_states["state1"], (batch_size_value, 1, 1)))
self.assertAllEqual(
state2_value,
np.tile(initial_states["state2"], (batch_size_value, 1)))
self.assertAllEqual(length_value, [2, 2])
# Step 2
(key_value, next_key_value, seq1_value, seq2_value, context1_value,
state1_value, state2_value, length_value, _, _) = sess.run(
(next_batch.key,
next_batch.next_key,
next_batch.sequences["seq1"],
next_batch.sequences["seq2"],
next_batch.context["context1"],
next_batch.state("state1"),
next_batch.state("state2"),
next_batch.length,
next_batch.save_state("state1", next_batch.state("state1") + 1),
next_batch.save_state("state2", next_batch.state("state2") - 1)))
self.assertEqual(set(key_value), expected_second_keys)
self.assertEqual(set(next_key_value), expected_final_keys)
self.assertAllEqual(
context1_value, np.tile(context["context1"], (batch_size_value, 1)))
self.assertAllEqual(
seq1_value, np.tile(sequences["seq1"][np.newaxis, 2:4, :],
(batch_size_value, 1, 1)))
self.assertAllEqual(
seq2_value,
np.tile(sequences["seq2"][np.newaxis, 2:4, :, :],
(batch_size_value, 1, 1, 1)))
self.assertAllEqual(
state1_value,
1 + np.tile(initial_states["state1"], (batch_size_value, 1, 1)))
self.assertAllEqual(
state2_value,
-1 + np.tile(initial_states["state2"], (batch_size_value, 1)))
self.assertAllEqual(length_value, [1, 1])
# Finished. Let's make sure there's nothing left in the barrier.
self.assertEqual(0, state_saver.barrier.ready_size().eval())
def main(_):
if not tf.gfile.Exists(FLAGS.train_log_dir):
tf.gfile.MakeDirs(FLAGS.train_log_dir)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Put input pipeline on CPU to reserve GPU for training.
with tf.name_scope('inputs'), tf.device('/cpu:0'):
images = data_provider.provide_data(
'train', FLAGS.batch_size, dataset_dir=FLAGS.dataset_dir,
patch_size=FLAGS.patch_size)
# Manually define a GANModel tuple. This is useful when we have custom
# code to track variables. Note that we could replace all of this with a
# call to `tfgan.gan_model`, but we don't in order to demonstrate some of
# TFGAN's flexibility.
with tf.variable_scope('generator') as gen_scope:
reconstructions, _, prebinary = networks.compression_model(
images,
num_bits=FLAGS.bits_per_patch,
depth=FLAGS.model_depth)
gan_model = _get_gan_model(
generator_inputs=images,
generated_data=reconstructions,
real_data=images,
generator_scope=gen_scope)
summaries.add_reconstruction_summaries(images, reconstructions, prebinary)
tfgan.eval.add_gan_model_summaries(gan_model)
# Define the GANLoss tuple using standard library functions.
with tf.name_scope('loss'):
gan_loss = tfgan.gan_loss(
gan_model,
generator_loss_fn=tfgan.losses.least_squares_generator_loss,
discriminator_loss_fn=tfgan.losses.least_squares_discriminator_loss,
add_summaries=FLAGS.weight_factor > 0)
# Define the standard pixel loss.
l1_pixel_loss = tf.norm(gan_model.real_data - gan_model.generated_data,
ord=1)
# Modify the loss tuple to include the pixel loss. Add summaries as well.
gan_loss = tfgan.losses.combine_adversarial_loss(
gan_loss, gan_model, l1_pixel_loss, weight_factor=FLAGS.weight_factor)
# Get the GANTrain ops using the custom optimizers and optional
# discriminator weight clipping.
with tf.name_scope('train_ops'):
gen_lr, dis_lr = _lr(FLAGS.generator_lr, FLAGS.discriminator_lr)
gen_opt, dis_opt = _optimizer(gen_lr, dis_lr)
train_ops = tfgan.gan_train_ops(
gan_model,
gan_loss,
generator_optimizer=gen_opt,
discriminator_optimizer=dis_opt,
summarize_gradients=True,
colocate_gradients_with_ops=True,
aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)
tf.summary.scalar('generator_lr', gen_lr)
tf.summary.scalar('discriminator_lr', dis_lr)
# Determine the number of generator vs discriminator steps.
train_steps = tfgan.GANTrainSteps(
generator_train_steps=1,
discriminator_train_steps=int(FLAGS.weight_factor > 0))
# Run the alternating training loop. Skip it if no steps should be taken
# (used for graph construction tests).
status_message = tf.string_join(
['Starting train step: ',
tf.as_string(tf.train.get_or_create_global_step())],
name='status_message')
if FLAGS.max_number_of_steps == 0: return
tfgan.gan_train(
train_ops,
FLAGS.train_log_dir,
tfgan.get_sequential_train_hooks(train_steps),
hooks=[tf.train.StopAtStepHook(num_steps=FLAGS.max_number_of_steps),
tf.train.LoggingTensorHook([status_message], every_n_iter=10)],
master=FLAGS.master,
is_chief=FLAGS.task == 0)
def main(_):
if not tf.gfile.Exists(FLAGS.train_log_dir):
tf.gfile.MakeDirs(FLAGS.train_log_dir)
# Force all input processing onto CPU in order to reserve the GPU for
# the forward inference and back-propagation.
with tf.name_scope('inputs'):
with tf.device('/cpu:0'):
images, one_hot_labels, _ = data_provider.provide_data(
'train', FLAGS.batch_size, FLAGS.dataset_dir, num_threads=4)
# Define the GANModel tuple. Optionally, condition the GAN on the label or
# use an InfoGAN to learn a latent representation.
if FLAGS.gan_type == 'unconditional':
gan_model = tfgan.gan_model(
generator_fn=networks.unconditional_generator,
discriminator_fn=networks.unconditional_discriminator,
real_data=images,
generator_inputs=tf.random_normal(
[FLAGS.batch_size, FLAGS.noise_dims]))
elif FLAGS.gan_type == 'conditional':
noise = tf.random_normal([FLAGS.batch_size, FLAGS.noise_dims])
gan_model = tfgan.gan_model(
generator_fn=networks.conditional_generator,
discriminator_fn=networks.conditional_discriminator,
real_data=images,
generator_inputs=(noise, one_hot_labels))
elif FLAGS.gan_type == 'infogan':
cat_dim, cont_dim = 10, 2
generator_fn = functools.partial(
networks.infogan_generator, categorical_dim=cat_dim)
discriminator_fn = functools.partial(
networks.infogan_discriminator, categorical_dim=cat_dim,
continuous_dim=cont_dim)
unstructured_inputs, structured_inputs = util.get_infogan_noise(
FLAGS.batch_size, cat_dim, cont_dim, FLAGS.noise_dims)
gan_model = tfgan.infogan_model(
generator_fn=generator_fn,
discriminator_fn=discriminator_fn,
real_data=images,
unstructured_generator_inputs=unstructured_inputs,
structured_generator_inputs=structured_inputs)
tfgan.eval.add_gan_model_image_summaries(gan_model, FLAGS.grid_size)
# Get the GANLoss tuple. You can pass a custom function, use one of the
# already-implemented losses from the losses library, or use the defaults.
with tf.name_scope('loss'):
mutual_information_penalty_weight = (1.0 if FLAGS.gan_type == 'infogan'
else 0.0)
gan_loss = tfgan.gan_loss(
gan_model,
gradient_penalty_weight=1.0,
mutual_information_penalty_weight=mutual_information_penalty_weight,
add_summaries=True)
tfgan.eval.add_regularization_loss_summaries(gan_model)
# Get the GANTrain ops using custom optimizers.
with tf.name_scope('train'):
gen_lr, dis_lr = _learning_rate(FLAGS.gan_type)
train_ops = tfgan.gan_train_ops(
gan_model,
gan_loss,
generator_optimizer=tf.train.AdamOptimizer(gen_lr, 0.5),
discriminator_optimizer=tf.train.AdamOptimizer(dis_lr, 0.5),
summarize_gradients=True,
aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)
# Run the alternating training loop. Skip it if no steps should be taken
# (used for graph construction tests).
status_message = tf.string_join(
['Starting train step: ',
tf.as_string(tf.train.get_or_create_global_step())],
name='status_message')
if FLAGS.max_number_of_steps == 0: return
tfgan.gan_train(
train_ops,
hooks=[tf.train.StopAtStepHook(num_steps=FLAGS.max_number_of_steps),
tf.train.LoggingTensorHook([status_message], every_n_iter=10)],
logdir=FLAGS.train_log_dir,
get_hooks_fn=tfgan.get_joint_train_hooks())
def main(_):
if not tf.gfile.Exists(FLAGS.train_log_dir):
tf.gfile.MakeDirs(FLAGS.train_log_dir)
with tf.device(tf.train.replica_device_setter(FLAGS.ps_tasks)):
# Force all input processing onto CPU in order to reserve the GPU for
# the forward inference and back-propagation.
with tf.name_scope('inputs'):
with tf.device('/cpu:0'):
images, one_hot_labels, _, _ = data_provider.provide_data(
FLAGS.batch_size, FLAGS.dataset_dir)
# Define the GANModel tuple.
noise = tf.random_normal([FLAGS.batch_size, 64])
if FLAGS.conditional:
generator_fn = networks.conditional_generator
discriminator_fn = networks.conditional_discriminator
generator_inputs = (noise, one_hot_labels)
else:
generator_fn = networks.generator
discriminator_fn = networks.discriminator
generator_inputs = noise
gan_model = tfgan.gan_model(
generator_fn,
discriminator_fn,
real_data=images,
generator_inputs=generator_inputs)
tfgan.eval.add_gan_model_image_summaries(gan_model)
# Get the GANLoss tuple. Use the selected GAN loss functions.
# (joelshor): Put this block in `with tf.name_scope('loss'):` when
# cl/171610946 goes into the opensource release.
gan_loss = tfgan.gan_loss(gan_model,
gradient_penalty_weight=1.0,
add_summaries=True)
# Get the GANTrain ops using the custom optimizers and optional
# discriminator weight clipping.
with tf.name_scope('train'):
gen_lr, dis_lr = _learning_rate()
gen_opt, dis_opt = _optimizer(gen_lr, dis_lr, FLAGS.use_sync_replicas)
train_ops = tfgan.gan_train_ops(
gan_model,
gan_loss,
generator_optimizer=gen_opt,
discriminator_optimizer=dis_opt,
summarize_gradients=True,
colocate_gradients_with_ops=True,
aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)
tf.summary.scalar('generator_lr', gen_lr)
tf.summary.scalar('discriminator_lr', dis_lr)
# Run the alternating training loop. Skip it if no steps should be taken
# (used for graph construction tests).
sync_hooks = ([gen_opt.make_session_run_hook(FLAGS.task == 0),
dis_opt.make_session_run_hook(FLAGS.task == 0)]
if FLAGS.use_sync_replicas else [])
status_message = tf.string_join(
['Starting train step: ',
tf.as_string(tf.train.get_or_create_global_step())],
name='status_message')
if FLAGS.max_number_of_steps == 0: return
sess_config = tf.ConfigProto(
inter_op_parallelism_threads=FLAGS.inter_op_parallelism_threads,
intra_op_parallelism_threads=FLAGS.intra_op_parallelism_threads)
tfgan.gan_train(
train_ops,
hooks=(
[tf.train.StopAtStepHook(num_steps=FLAGS.max_number_of_steps),
tf.train.LoggingTensorHook([status_message], every_n_iter=10)] +
sync_hooks),
logdir=FLAGS.train_log_dir,
master=FLAGS.master,
is_chief=FLAGS.task == 0,
config=sess_config)
