def init(self, reset=True):
if self.sess is None:
print 'Running on:', self.gpu
with tf.device(self.gpu):
if reset:
tf.reset_default_graph()
tf.Graph().as_default()
pr = self.pr
self.sess = tf.Session()
self.ims_ph = tf.placeholder(
tf.uint8,
[1, pr.sampled_frames, pr.crop_im_dim, pr.crop_im_dim, 3])
self.samples_ph = tf.placeholder(tf.float32,
(1, pr.num_samples, 2))
crop_spec = lambda x: x[:, :pr.spec_len]
samples_trunc = self.samples_ph[:, :pr.sample_len]
spec_mix, phase_mix = sourcesep.stft(samples_trunc[:, :, 0],
pr)
print 'Raw spec length:', mu.shape(spec_mix)
spec_mix = crop_spec(spec_mix)
phase_mix = crop_spec(phase_mix)
print 'Truncated spec length:', mu.shape(spec_mix)
self.specgram_op, phase = map(
crop_spec, sourcesep.stft(samples_trunc[:, :, 0], pr))
self.auto_op = sourcesep.istft(self.specgram_op, phase, pr)
self.net = sourcesep.make_net(self.ims_ph,
samples_trunc,
spec_mix,
phase_mix,
pr,
reuse=False,
train=False)
self.spec_pred_fg = self.net.pred_spec_fg
self.spec_pred_bg = self.net.pred_spec_bg
self.samples_pred_fg = self.net.pred_wav_fg
self.samples_pred_bg = self.net.pred_wav_bg
print 'Restoring from:', pr.model_path
if self.restore_only_shift:
print 'restoring only shift'
import tensorflow.contrib.slim as slim
var_list = slim.get_variables_to_restore()
var_list = [
x for x in var_list
if x.name.startswith('im/') or x.name.startswith('sf/')
or x.name.startswith('joint/')
]
self.sess.run(tf.global_variables_initializer())
tf.train.Saver(var_list).restore(self.sess, pr.model_path)
else:
tf.train.Saver().restore(self.sess, pr.model_path)
tf.get_default_graph().finalize()
def init(self, reset=True):
if self.sess is None:
print 'Running on:', self.gpu
with tf.device(self.gpu):
if reset:
tf.reset_default_graph()
tf.Graph().as_default()
pr = self.pr
self.sess = tf.Session()
self.ims_ph = tf.placeholder(
tf.uint8,
[1, pr.sampled_frames, pr.crop_im_dim, pr.crop_im_dim, 3])
self.samples_ph = tf.placeholder(tf.float32,
(1, pr.num_samples, 2))
crop_spec = lambda x: x[:, :pr.spec_len]
samples_trunc = self.samples_ph[:, :pr.sample_len]
spec_mix, phase_mix = sourcesep.stft(samples_trunc[:, :, 0],
pr)
print 'Raw spec length:', mu.shape(spec_mix)
spec_mix = crop_spec(spec_mix)
phase_mix = crop_spec(phase_mix)
print 'Truncated spec length:', mu.shape(spec_mix)
self.specgram_op, phase = map(
crop_spec, sourcesep.stft(samples_trunc[:, :, 0], pr))
self.auto_op = sourcesep.istft(self.specgram_op, phase, pr)
self.net = sourcesep.make_net(self.ims_ph,
samples_trunc,
spec_mix,
phase_mix,
pr,
reuse=False,
train=False)
self.spec_pred_fg = self.net.pred_spec_fg
self.spec_pred_bg = self.net.pred_spec_bg
self.samples_pred_fg = self.net.pred_wav_fg
self.samples_pred_bg = self.net.pred_wav_bg
print 'Restoring from:', pr.model_path
if self.restore_only_shift:
print 'restoring only shift'
import tensorflow.contrib.slim as slim
var_list = slim.get_variables_to_restore()
var_list = [
x for x in var_list
if x.name.startswith('im/') or x.name.startswith('sf/')
or x.name.startswith('joint/')
]
self.sess.run(tf.global_variables_initializer())
tf.train.Saver(var_list).restore(self.sess, pr.model_path)
else:
tf.train.Saver().restore(self.sess, pr.model_path)
# NEW STUFF!
# inputs = {'input': tf.saved_model.utils.build_tensor_info(ims_ph)}
# out_classes = sess.graph.get_tensor_by_name('AttentionOcr_v1/ReduceJoin:0')
# outputs = {'output': tf.saved_model.utils.build_tensor_info(out_classes)}
signature = tf.saved_model.signature_def_utils.build_signature_def(
inputs={
"samples_ph":
tf.saved_model.utils.build_tensor_info(
self.samples_ph),
"ims_ph":
tf.saved_model.utils.build_tensor_info(self.ims_ph),
},
outputs={
"spec_pred_fg":
tf.saved_model.utils.build_tensor_info(
self.spec_pred_fg),
"spec_pred_bg":
tf.saved_model.utils.build_tensor_info(
self.spec_pred_bg),
"samples_pred_fg":
tf.saved_model.utils.build_tensor_info(
self.samples_pred_fg),
"samples_pred_bg":
tf.saved_model.utils.build_tensor_info(
self.samples_pred_bg),
"specgram_op":
tf.saved_model.utils.build_tensor_info(
self.specgram_op)
},
method_name=tf.saved_model.signature_constants.
PREDICT_METHOD_NAME)
# legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
print("+++++++ EXPORTING MODEL +++++++")
builder = tf.saved_model.builder.SavedModelBuilder(
"../exported-1/")
builder.add_meta_graph_and_variables(
self.sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map={
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
signature
}) #,
#legacy_init_op=legacy_init_op)
builder.save()
print("+++++++ DONE EXPORTING MODEL +++++++")
tf.get_default_graph().finalize()