catch SIGTERM

show messages on index

Author: xymeow

Diff for: img_prep.py

@@ -28,12 +28,12 @@ def deprocess_image(x):
     return x
 
 
-def random_image():
-    """
-    Create a random image
-    :return: A random image
-    """
-    return np.random.uniform(0, 255, (1, 3, img_width, img_height))
+# def random_image():
+#     """
+#     Create a random image
+#     :return: A random image
+#     """
+#     return np.random.uniform(0, 255, (1, 3, img_width, img_height))
 
 
 def grey_image():

Diff for: index.html

@@ -40,22 +40,22 @@
         if (trained){
             console.log('pause');
             $('#train').get(0).innerHTML = 'train';
-            $('#status').get(0).innerHTML = 'pause';
+//            $('#status').get(0).innerHTML = 'pause';
             trained = false;
         }
 
         else {
             console.log('train');
             $('#train').get(0).innerHTML = 'pause';
-            $('#status').get(0).innerHTML = 'training';
+//            $('#status').get(0).innerHTML = 'training';
             trained = true;
         }
     }
 
     function stop() {
         $.ajax('/stop')
         $('#train').get(0).innerHTML = 'train';
-        $('#status').get(0).innerHTML = 'stop';
+//        $('#status').get(0).innerHTML = 'stop';
     }
 
     $(document).ready(function() {
@@ -64,14 +64,16 @@
             $.post('http://localhost:8000/lr?lr='+this.value);
     });
         updater.poll();
+//        alert(updater2);
+        updater2.poll();
   });
 
 
 
 </script>
 <!--<span id="test" style="padding-left: 240px;"></span>-->
 <ul class="side-nav fixed">
-    <li style="text-align: center"><b >status: <span id="status">stop</span></b></li>
+    <li style="text-align: center"><b >status:</b> <br><span id="status">stop</span></li>
     <li><a class="waves-effect waves-light btn cyan accent-4 white-text" href="javascript: train()" id="train">train</a></li>
     <li><a class="waves-effect waves-light btn cyan accent-4 white-text" href="javascript: stop()" id="stop">stop</a></li>
     <li><a class="waves-effect waves-light btn cyan accent-4 white-text" href="javascript: toDefault()">Default View</a></li>

Diff for: neural_style.py

@@ -1,30 +1,27 @@
 from __future__ import print_function
 import os
-import sys
-
-import multiprocessing
-from start_training import start_training
-
-import json
-import uuid
 from tornado.ioloop import IOLoop
 from tornado.web import RequestHandler, Application
 from tornado.escape import json_decode
+import json
+# from main import start_training
+import multiprocessing
 from tornado.concurrent import Future
 from tornado import gen
+import uuid
+# from training_record import TrainingRecorder
 
-__author__ = 'xymeow', 'suquark'
+__author__ = 'xymeow'
 
 # TODO: show training messages on the front end
 
-label = None
-
 start = False
 trained = False
 stop = False
 lr = 1.0
 
-
+p = None
+record_cache = None
 # use long polling to update the training results
 # I copied these codes from tornado/demo/chat
 class MessageBuffer(object):
@@ -67,7 +64,7 @@ def new_messages(self, messages):
 
 
 record_pool = MessageBuffer()
-
+message_pool = MessageBuffer()
 
 class MainHandler(RequestHandler):
     def get(self):
@@ -79,10 +76,13 @@ def post(self):
 
 
 class RecordHandler(RequestHandler):
+
     def post(self):
+        global record_cache
         json_dict = json_decode(self.request.body)
         print(self.request.body)
-
+        record_cache = json_dict['cache']
+        json_dict.pop('cache')
         myjson = {
             'id': str(uuid.uuid4()),
             'json': json_dict
@@ -112,11 +112,14 @@ def on_connection_close(self):
 
 class TrainHandler(RequestHandler):
     def get(self):
-        global trained, start, label
+        global trained
+        global start
+        global p
         if not start:
             start = True
             trained = True
-            multiprocessing.Process(target=start_training).start()
+            p = multiprocessing.Process(target=start_training)
+            p.start()
         elif trained:
             trained = False
         elif not trained:
@@ -139,8 +142,11 @@ class StopHandler(RequestHandler):
     def get(self):
         global stop
         global start
+        global p
         stop = True
         start = False
+        if p != None:
+            p.terminate()
 
 
 class PictureHandler(RequestHandler):
@@ -153,12 +159,12 @@ def get(self):
 
 class InitJSONHandler(RequestHandler):
     def get(self):
-        if os.path.exists(label + '.json'):
-            with open(label + '.json', 'rb') as f:
+        global record_cache
+        if record_cache == None:
+            with open('result.json', 'rb') as f:
                 self.write(f.read())
         else:
-            print('Record file not found.')
-            self.write('{}')
+            self.write(record_cache)
 
 
 class LearningRateHandler(RequestHandler):
@@ -173,6 +179,36 @@ def post(self):
         print('set learnig rate to {}'.format(lr))
 
 
+class PullMessageHandler(RequestHandler):
+
+    @gen.coroutine
+    def post(self):
+        cursor = self.get_argument("cursor", None)
+        # Save the future returned by wait_for_messages so we can cancel
+        # it in wait_for_messages
+        self.future = message_pool.wait_for_messages(cursor=cursor)
+        messages = yield self.future
+        if self.request.connection.stream.closed():
+            return
+        self.write(dict(messages=messages))
+
+    def on_connection_close(self):
+        message_pool.cancel_wait(self.future)
+
+
+class SetMessageHandler(RequestHandler):
+    def post(self):
+        # global record_cache
+        json_dict = json_decode(self.request.body)
+        print(self.request.body)
+        myjson = {
+            'id': str(uuid.uuid4()),
+            'json': json_dict
+        }
+        message_pool.new_messages([myjson])
+
+
+
 settings = {
     "static_path": os.path.join(os.path.dirname(__file__), "static")
 }
@@ -186,11 +222,12 @@ def post(self):
     (r"/picture", PictureHandler),
     (r"/init", InitJSONHandler),
     (r"/lr", LearningRateHandler),
-    (r"/train", TrainHandler)
+    (r"/train", TrainHandler),
+    (r"/setmsg", SetMessageHandler),
+    (r"/getmsg", PullMessageHandler)
 ], **settings)
 
 if __name__ == "__main__":
-    # assert len(sys.argv) > 1, 'You should give the name of this training'
-    label = sys.argv[1] if len(sys.argv) > 1 else 'result'  # Which record would you like to take?
+    print('server start at localhost:8000')
     application.listen(8000)
     IOLoop.instance().start()

Diff for: start_training.py

@@ -9,20 +9,43 @@
 from training_record import TrainingRecorder
 from img_prep import img_in, img_save, preprocess_image, grey_image
 from settings import content_path, style_path, loss_set, result_dir
-import requests, time
+import requests, time, signal, sys
 
 
 # load content and style image
+# record_cache = None
+tr = None
+
+
+def signal_term_handler(signal, frame):
+    global tr
+    if tr == None:
+        print('stop training.')
+        requests.post('http://localhost:8000/setmsg', None,
+                              {'msg': 'stop training.'})
+        # sys.exit(0)
+    else:
+        print('stop training and save weights into json file.')
+        requests.post('http://localhost:8000/setmsg', None,
+                              {'msg': 'stop training and save weights into json file.'})
+        tr.export_file()
+    sys.exit(0)
+
+signal.signal(signal.SIGTERM, signal_term_handler)
+
+
 def start_training():
-    content, style = img_in(content_path, style_path)
 
-    tr = TrainingRecorder(result_dir, loss_set, img_save, '.png')
+    content, style = img_in(content_path, style_path)
+    global tr
+    tr = TrainingRecorder(loss_set, result_dir, img_save, '.png')
 
-    idx = tr.get_tail() - 1
-    if idx >= 0:
+    idx = tr.get_head() - 1
+    if idx > 0:
         x = preprocess_image(tr.get_name(idx))
     else:
-        x = grey_image()
+        x = content
+
     idx += 1
 
     model = Model(content, style, x)
@@ -34,13 +57,16 @@ def start_training():
             lr = float(requests.get('http://localhost:8000/lr').text)
             if status == 'pause':
                 i -= 1
+                requests.post('http://localhost:8000/setmsg', None,
+                              {'msg': 'pause'})
                 while True:
                     time.sleep(1)
                     status = requests.get('http://localhost:8000/status').text
                     if status != 'pause':
                         break
             elif status == 'training':
                 print('Start of iteration', i)
+                requests.post('http://localhost:8000/setmsg', None, {'msg': 'Start of iteration ' + str(i)})
                 # msg = 'Start of iteration ' + str(i)
                 if lst_lr != lr:
                     # if learning rate has changes, set a new lr for the optimizer
@@ -51,9 +77,14 @@ def start_training():
                 lst_lr = lr
             elif status == 'stop':
                 print('stop training and save weights into json file.')
+                requests.post('http://localhost:8000/setmsg', None,
+                              {'msg': 'stop training and save weights into json file.'})
                 # msg = 'stop training and save weights into json file.'
                 tr.export_file()
                 break
 
     except KeyboardInterrupt:
+        print('stop training and save weights into json file.')
+        requests.post('http://localhost:8000/setmsg', None,
+                              {'msg': 'stop training and save weights into json file.'})
         tr.export_file()

Diff for: training.py

@@ -4,6 +4,7 @@
 from settings import img_height, img_width, input_shape, learning_rate
 from settings import loss_set, content_feature_layers, style_feature_layers
 from vgg16_cnnonly_model import set_model_input, model_with_input
+import requests
 
 
 def gram_matrix(x):
@@ -39,6 +40,8 @@ def __init__(self, content, style, x):
         self.model = None
         # evaluate features of content and style images
         print("Pre-evaluate features...")
+        requests.post('http://localhost:8000/setmsg', None,
+                              {'msg': "Pre-evaluate features..."})
         # global msg
         # msg = "Pre-evaluate features..."
         self.gen_model(K.placeholder(input_shape))
@@ -60,6 +63,8 @@ def set_lr(self, learning_rate):
         # self.optimizer.lr.set_value(learning_rate)
         K.set_value(self.optimizer.lr, learning_rate)
         print('learning rate = {}'.format(learning_rate))
+        requests.post('http://localhost:8000/setmsg', None,
+                              {'msg': 'set learning rate to {}'.format(learning_rate)})
         # global msg
         # msg = 'learning rate = {}'.format(learning_rate)
 
@@ -146,6 +151,8 @@ def compile(self):
         """
         # global msg
         print("Generate loss and grad...")
+        requests.post('http://localhost:8000/setmsg', None,
+                              {'msg': "Generate loss and grad..."})
         # msg = "Generate loss and grad..."
         losses = [l * w for l, w in zip(*self.get_loss())]
         total_loss = sum(losses)