Merge pull request #47 from ChunweiYan/feature/refactor_python_sdk

server/visualdl/lib.py

@@ -16,47 +16,47 @@ def get_modes(storage):
 def get_scalar_tags(storage, mode):
     result = {}
     for mode in storage.modes():
-        reader = storage.as_mode(mode)
-        tags = reader.tags('scalar')
-        if tags:
-            result[mode] = {}
-            for tag in tags:
-                result[mode][tag] = {
-                    'displayName': reader.scalar(tag).caption(),
-                    'description': "",
-                }
+        with storage.mode(mode) as reader:
+            tags = reader.tags('scalar')
+            if tags:
+                result[mode] = {}
+                for tag in tags:
+                    result[mode][tag] = {
+                        'displayName': reader.scalar(tag).caption(),
+                        'description': "",
+                    }
     return result
 
 
 def get_scalar(storage, mode, tag):
-    reader = storage.as_mode(mode)
-    scalar = reader.scalar(tag)
+    with storage.mode(mode) as reader:
+        scalar = reader.scalar(tag)
 
-    records = scalar.records()
-    ids = scalar.ids()
-    timestamps = scalar.timestamps()
+        records = scalar.records()
+        ids = scalar.ids()
+        timestamps = scalar.timestamps()
 
-    result = zip(timestamps, ids, records)
-    return result
+        result = zip(timestamps, ids, records)
+        return result
 
 
 def get_image_tags(storage):
     result = {}
 
     for mode in storage.modes():
-        reader = storage.as_mode(mode)
-        tags = reader.tags('image')
-        if tags:
-            result[mode] = {}
-            for tag in tags:
-                image = reader.image(tag)
-                for i in xrange(max(1, image.num_samples())):
-                    caption = tag if image.num_samples() <= 1 else '%s/%d'%(tag, i)
-                    result[mode][caption] = {
-                        'displayName': caption,
-                        'description': "",
-                        'samples': 1,
-                    }
+        with storage.mode(mode) as reader:
+            tags = reader.tags('image')
+            if tags:
+                result[mode] = {}
+                for tag in tags:
+                    image = reader.image(tag)
+                    for i in xrange(max(1, image.num_samples())):
+                        caption = tag if image.num_samples() <= 1 else '%s/%d'%(tag, i)
+                        result[mode][caption] = {
+                            'displayName': caption,
+                            'description': "",
+                            'samples': 1,
+                        }
     return result
 
 
@@ -70,9 +70,9 @@ def get_image_tag_steps(storage, mode, tag):
         tag = tag[:tag.rfind('/')]
         sample_index = int(res.groups()[0])
 
-    reader = storage.as_mode(mode)
-    image = reader.image(tag)
-    res = []
+    with storage.mode(mode) as reader:
+        image = reader.image(tag)
+        res = []
 
     for step_index in range(image.num_records()):
         record = image.record(step_index, sample_index)
@@ -95,26 +95,26 @@ def get_image_tag_steps(storage, mode, tag):
 
 
 def get_invididual_image(storage, mode, tag, step_index):
-    reader = storage.as_mode(mode)
-    res = re.search(r".*/([0-9]+$)", tag)
-    # remove suffix '/x'
-    if res:
-        offset = int(res.groups()[0])
-        tag = tag[:tag.rfind('/')]
+    with storage.mode(mode) as reader:
+        res = re.search(r".*/([0-9]+$)", tag)
+        # remove suffix '/x'
+        if res:
+            offset = int(res.groups()[0])
+            tag = tag[:tag.rfind('/')]
 
-    image = reader.image(tag)
-    record = image.record(step_index, offset)
+        image = reader.image(tag)
+        record = image.record(step_index, offset)
 
-    data = np.array(record.data(), dtype='uint8').reshape(record.shape())
-    tempfile = NamedTemporaryFile(mode='w+b', suffix='.png')
-    with Image.fromarray(data) as im:
-        im.save(tempfile)
-    tempfile.seek(0, 0)
-    return tempfile
+        data = np.array(record.data(), dtype='uint8').reshape(record.shape())
+        tempfile = NamedTemporaryFile(mode='w+b', suffix='.png')
+        with Image.fromarray(data) as im:
+            im.save(tempfile)
+        tempfile.seek(0, 0)
+        return tempfile
 
 
 if __name__ == '__main__':
-    reader = storage.StorageReader('./tmp/mock')
+    reader = storage.LogReader('./tmp/mock')
     tags = get_image_tags(reader)
 
     tags = get_image_tag_steps(reader, 'train', 'layer1/layer2/image0/0')

server/visualdl/lib_test.py

@@ -8,7 +8,7 @@ from storage_mock import add_scalar, add_image
 class LibTest(unittest.TestCase):
     def setUp(self):
         dir = "./tmp/mock"
-        writer = storage.StorageWriter(dir, sync_cycle=20)
+        writer = storage.LogWriter(dir, sync_cycle=20)
 
         add_scalar(writer, "train", "layer/scalar0/min", 1000, 1)
         add_scalar(writer, "test", "layer/scalar0/min", 1000, 10)
@@ -24,7 +24,7 @@ class LibTest(unittest.TestCase):
         add_image(writer, "train", "layer/image1", 7, 10, 1, shape=[30,30,2])
         add_image(writer, "test", "layer/image1", 7, 10, 1, shape=[30,30,2])
 
-        self.reader = storage.StorageReader(dir)
+        self.reader = storage.LogReader(dir)
 
     def test_modes(self):
         modes = lib.get_modes(self.reader)

server/visualdl/storage_mock.py

@@ -6,11 +6,11 @@ import numpy as np
 
 
 def add_scalar(writer, mode, tag, num_steps, skip):
-    my_writer = writer.as_mode(mode)
-    scalar = my_writer.scalar(tag)
-    for i in range(num_steps):
-        if i % skip == 0:
-            scalar.add_record(i, random.random())
+    with writer.mode(mode) as my_writer:
+        scalar = my_writer.scalar(tag)
+        for i in range(num_steps):
+            if i % skip == 0:
+                scalar.add_record(i, random.random())
 
 
 def add_image(writer,
@@ -20,20 +20,20 @@ def add_image(writer,
               num_passes,
               step_cycle,
               shape=[50, 50, 3]):
-    writer_ = writer.as_mode(mode)
-    image_writer = writer_.image(tag, num_samples, step_cycle)
-
-    for pass_ in xrange(num_passes):
-        image_writer.start_sampling()
-        for ins in xrange(2 * num_samples):
-            index = image_writer.is_sample_taken()
-            if index != -1:
-                data = np.random.random(shape) * 256
-                data = np.ndarray.flatten(data)
-                assert shape
-                assert len(data) > 0
-                image_writer.set_sample(index, shape, list(data))
-        image_writer.finish_sampling()
+    with writer.mode(mode) as writer_:
+        image_writer = writer_.image(tag, num_samples, step_cycle)
+
+        for pass_ in xrange(num_passes):
+            image_writer.start_sampling()
+            for ins in xrange(2 * num_samples):
+                index = image_writer.is_sample_taken()
+                if index != -1:
+                    data = np.random.random(shape) * 256
+                    data = np.ndarray.flatten(data)
+                    assert shape
+                    assert len(data) > 0
+                    image_writer.set_sample(index, shape, list(data))
+            image_writer.finish_sampling()
 
 
 if __name__ == '__main__':

visualdl/logic/pybind.cc

@@ -11,49 +11,49 @@ namespace cp = visualdl::components;
 PYBIND11_PLUGIN(core) {
   py::module m("core", "C++ core of VisualDL");
 
-#define READER_ADD_SCALAR(T)                                            \
-  .def("get_scalar_" #T, [](vs::Reader& self, const std::string& tag) { \
-    auto tablet = self.tablet(tag);                                     \
-    return vs::components::ScalarReader<T>(std::move(tablet));          \
+#define READER_ADD_SCALAR(T)                                               \
+  .def("get_scalar_" #T, [](vs::LogReader& self, const std::string& tag) { \
+    auto tablet = self.tablet(tag);                                        \
+    return vs::components::ScalarReader<T>(std::move(tablet));             \
   })
-  py::class_<vs::Reader>(m, "Reader")
+  py::class_<vs::LogReader>(m, "LogReader")
       .def("__init__",
-           [](vs::Reader& instance, const std::string& dir) {
-             new (&instance) vs::Reader(dir);
+           [](vs::LogReader& instance, const std::string& dir) {
+             new (&instance) vs::LogReader(dir);
            })
-      .def("as_mode", &vs::Reader::AsMode)
-      .def("modes", [](vs::Reader& self) { return self.storage().modes(); })
-      .def("tags", &vs::Reader::tags)
+      .def("as_mode", &vs::LogReader::AsMode)
+      .def("modes", [](vs::LogReader& self) { return self.storage().modes(); })
+      .def("tags", &vs::LogReader::tags)
       // clang-format off
       READER_ADD_SCALAR(float)
       READER_ADD_SCALAR(double)
       READER_ADD_SCALAR(int)
       // clang-format on
-      .def("get_image", [](vs::Reader& self, const std::string& tag) {
+      .def("get_image", [](vs::LogReader& self, const std::string& tag) {
         auto tablet = self.tablet(tag);
         return vs::components::ImageReader(self.mode(), tablet);
       });
 #undef READER_ADD_SCALAR
 
-#define WRITER_ADD_SCALAR(T)                                            \
-  .def("new_scalar_" #T, [](vs::Writer& self, const std::string& tag) { \
-    auto tablet = self.AddTablet(tag);                                  \
-    return cp::Scalar<T>(tablet);                                       \
+#define WRITER_ADD_SCALAR(T)                                               \
+  .def("new_scalar_" #T, [](vs::LogWriter& self, const std::string& tag) { \
+    auto tablet = self.AddTablet(tag);                                     \
+    return cp::Scalar<T>(tablet);                                          \
   })
 
-  py::class_<vs::Writer>(m, "Writer")
+  py::class_<vs::LogWriter>(m, "LogWriter")
       .def("__init__",
-           [](vs::Writer& instance, const std::string& dir, int sync_cycle) {
-             new (&instance) vs::Writer(dir, sync_cycle);
+           [](vs::LogWriter& instance, const std::string& dir, int sync_cycle) {
+             new (&instance) vs::LogWriter(dir, sync_cycle);
            })
-      .def("as_mode", &vs::Writer::AsMode)
+      .def("as_mode", &vs::LogWriter::AsMode)
       // clang-format off
       WRITER_ADD_SCALAR(float)
       WRITER_ADD_SCALAR(double)
       WRITER_ADD_SCALAR(int)
       // clang-format on
       .def("new_image",
-           [](vs::Writer& self,
+           [](vs::LogWriter& self,
               const std::string& tag,
               int num_samples,
               int step_cycle) {
@@ -93,10 +93,7 @@ PYBIND11_PLUGIN(core) {
 
   py::class_<cp::ImageReader::ImageRecord>(m, "ImageRecord")
       // TODO(ChunweiYan) make these copyless.
-      .def("data",
-           [](cp::ImageReader::ImageRecord& self) {
-             return self.data;
-           })
+      .def("data", [](cp::ImageReader::ImageRecord& self) { return self.data; })
       .def("shape",
            [](cp::ImageReader::ImageRecord& self) { return self.shape; })
       .def("step_id",

visualdl/logic/sdk.cc

@@ -139,8 +139,8 @@ void Image::SetSample(int index,
 std::string ImageReader::caption() {
   CHECK_EQ(reader_.captions().size(), 1);
   auto caption = reader_.captions().front();
-  if (Reader::TagMatchMode(caption, mode_)) {
-    return Reader::GenReadableTag(mode_, caption);
+  if (LogReader::TagMatchMode(caption, mode_)) {
+    return LogReader::GenReadableTag(mode_, caption);
   }
   string::TagDecode(caption);
   return caption;

visualdl/logic/sdk.h

@@ -8,19 +8,19 @@ namespace visualdl {
 
 const static std::string kDefaultMode{"default"};
 
-class Writer {
+class LogWriter {
 public:
-  Writer(const std::string& dir, int sync_cycle) {
+  LogWriter(const std::string& dir, int sync_cycle) {
     storage_.SetDir(dir);
     storage_.meta.cycle = sync_cycle;
   }
-  Writer(const Writer& other) {
+  LogWriter(const LogWriter& other) {
     storage_ = other.storage_;
     mode_ = other.mode_;
   }
 
-  Writer AsMode(const std::string& mode) {
-    Writer writer = *this;
+  LogWriter AsMode(const std::string& mode) {
+    LogWriter writer = *this;
     storage_.AddMode(mode);
     writer.mode_ = mode;
     return writer;
@@ -43,11 +43,11 @@ private:
   std::string mode_{kDefaultMode};
 };
 
-class Reader {
+class LogReader {
 public:
-  Reader(const std::string& dir) : reader_(dir) {}
+  LogReader(const std::string& dir) : reader_(dir) {}
 
-  Reader AsMode(const std::string& mode) {
+  LogReader AsMode(const std::string& mode) {
     auto tmp = *this;
     tmp.mode_ = mode;
     return tmp;

visualdl/logic/sdk_test.cc

@@ -8,7 +8,7 @@ namespace visualdl {
 
 TEST(Scalar, write) {
   const auto dir = "./tmp/sdk_test";
-  Writer writer__(dir, 1);
+  LogWriter writer__(dir, 1);
   auto writer = writer__.AsMode("train");
   // write disk every time
   auto tablet = writer.AddTablet("scalar0");
@@ -19,7 +19,7 @@ TEST(Scalar, write) {
   scalar1.SetCaption("customized caption");
 
   // read from disk
-  Reader reader_(dir);
+  LogReader reader_(dir);
   auto reader = reader_.AsMode("train");
   auto tablet_reader = reader.tablet("scalar0");
   auto scalar_reader = components::ScalarReader<int>(std::move(tablet_reader));
@@ -44,7 +44,7 @@ TEST(Scalar, write) {
 
 TEST(Image, test) {
   const auto dir = "./tmp/sdk_test.image";
-  Writer writer__(dir, 4);
+  LogWriter writer__(dir, 4);
   auto writer = writer__.AsMode("train");
 
   auto tablet = writer.AddTablet("image0");
@@ -71,7 +71,7 @@ TEST(Image, test) {
 
   LOG(INFO) << "read images";
   // read it
-  Reader reader__(dir);
+  LogReader reader__(dir);
   auto reader = reader__.AsMode("train");
   auto tablet2read = reader.tablet("image0");
   components::ImageReader image2read("train", tablet2read);

visualdl/python/__init__.py

+from storage import *

visualdl/python/storage.py

@@ -7,14 +7,20 @@ import core
 dtypes = ("float", "double", "int32", "int64")
 
 
-class StorageReader(object):
+class LogReader(object):
+
+    cur_mode = None
 
     def __init__(self, dir, reader=None):
         self.dir = dir
-        self.reader = reader if reader else core.Reader(dir)
+        self.reader = reader if reader else core.LogReader(dir)
+
+    def mode(self, mode):
+        LogReader.cur_mode = self.as_mode(mode)
+        return LogReader.cur_mode
 
     def as_mode(self, mode):
-        tmp = StorageReader(dir, self.reader.as_mode(mode))
+        tmp = LogReader(dir, self.reader.as_mode(mode))
         return tmp
 
     def modes(self):
@@ -34,17 +40,29 @@ class StorageReader(object):
     def image(self, tag):
         return self.reader.get_image(tag)
 
+    def __enter__(self):
+        return LogReader.cur_mode
+
+    def __exit__(self, type, value, traceback):
+        pass
+
+
+class LogWriter(object):
 
-class StorageWriter(object):
+    cur_mode = None
 
     def __init__(self, dir, sync_cycle, writer=None):
         self.dir = dir
         self.sync_cycle = sync_cycle
-        self.writer = writer if writer else core.Writer(dir, sync_cycle)
+        self.writer = writer if writer else core.LogWriter(dir, sync_cycle)
+
+    def mode(self, mode):
+        LogWriter.cur_mode = self.as_mode(mode)
+        return LogWriter.cur_mode
 
     def as_mode(self, mode):
-        tmp = StorageWriter(self.dir, self.sync_cycle, self.writer.as_mode(mode))
-        return tmp
+        LogWriter.cur_mode = LogWriter(self.dir, self.sync_cycle, self.writer.as_mode(mode))
+        return LogWriter.cur_mode
 
     def scalar(self, tag, type='float'):
         type2scalar = {
@@ -56,3 +74,9 @@ class StorageWriter(object):
 
     def image(self, tag, num_samples, step_cycle):
         return self.writer.new_image(tag, num_samples, step_cycle)
+
+    def __enter__(self):
+        return LogWriter.cur_mode
+
+    def __exit__(self, type, value, traceback):
+        pass

visualdl/python/test_storage.py

@@ -11,7 +11,7 @@ import storage
 class StorageTest(unittest.TestCase):
     def setUp(self):
         self.dir = "./tmp/storage_test"
-        self.writer = storage.StorageWriter(
+        self.writer = storage.LogWriter(
             self.dir, sync_cycle=1).as_mode("train")
 
     def test_scalar(self):
@@ -22,15 +22,16 @@ class StorageTest(unittest.TestCase):
             scalar.add_record(i, float(i))
 
         print 'test read'
-        self.reader = storage.StorageReader(self.dir).as_mode("train")
-        scalar = self.reader.scalar("model/scalar/min")
-        self.assertEqual(scalar.caption(), "train")
-        records = scalar.records()
-        ids = scalar.ids()
-        self.assertTrue(np.equal(records, [float(i) for i in range(10)]).all())
-        self.assertTrue(np.equal(ids, [float(i) for i in range(10)]).all())
-        print 'records', records
-        print 'ids', ids
+        self.reader = storage.LogReader(self.dir)
+        with self.reader.mode("train") as reader:
+            scalar = reader.scalar("model/scalar/min")
+            self.assertEqual(scalar.caption(), "train")
+            records = scalar.records()
+            ids = scalar.ids()
+            self.assertTrue(np.equal(records, [float(i) for i in range(10)]).all())
+            self.assertTrue(np.equal(ids, [float(i) for i in range(10)]).all())
+            print 'records', records
+            print 'ids', ids
 
     def test_image(self):
         tag = "layer1/layer2/image0"
@@ -49,19 +50,20 @@ class StorageTest(unittest.TestCase):
                     image_writer.set_sample(index, shape, list(data))
             image_writer.finish_sampling()
 
-        self.reader = storage.StorageReader(self.dir).as_mode("train")
-        image_reader = self.reader.image(tag)
-        self.assertEqual(image_reader.caption(), tag)
-        self.assertEqual(image_reader.num_records(), num_passes)
+        self.reader = storage.LogReader(self.dir)
+        with self.reader.mode("train") as reader:
+            image_reader = reader.image(tag)
+            self.assertEqual(image_reader.caption(), tag)
+            self.assertEqual(image_reader.num_records(), num_passes)
 
-        image_record = image_reader.record(0, 1)
-        self.assertTrue(np.equal(image_record.shape(), shape).all())
-        data = image_record.data()
-        self.assertEqual(len(data), np.prod(shape))
+            image_record = image_reader.record(0, 1)
+            self.assertTrue(np.equal(image_record.shape(), shape).all())
+            data = image_record.data()
+            self.assertEqual(len(data), np.prod(shape))
 
-        image_tags = self.reader.tags("image")
-        self.assertTrue(image_tags)
-        self.assertEqual(len(image_tags), 1)
+            image_tags = reader.tags("image")
+            self.assertTrue(image_tags)
+            self.assertEqual(len(image_tags), 1)
 
     def test_check_image(self):
         '''
@@ -75,31 +77,44 @@ class StorageTest(unittest.TestCase):
         shape = [image.size[1], image.size[0], 3]
         origin_data = np.array(image.getdata()).flatten()
 
-        self.reader = storage.StorageReader(self.dir).as_mode("train")
+        self.reader = storage.LogReader(self.dir)
+        with self.reader.mode("train") as reader:
 
-        image_writer.start_sampling()
-        index = image_writer.is_sample_taken()
-        image_writer.set_sample(index, shape, list(origin_data))
-        image_writer.finish_sampling()
+            image_writer.start_sampling()
+            index = image_writer.is_sample_taken()
+            image_writer.set_sample(index, shape, list(origin_data))
+            image_writer.finish_sampling()
+
+            # read and check whether the original image will be displayed
+
+            image_reader = reader.image(tag)
+            image_record = image_reader.record(0, 0)
+            data = image_record.data()
+            shape = image_record.shape()
+
+            PIL_image_shape = (shape[0] * shape[1], shape[2])
+            data = np.array(data, dtype='uint8').reshape(PIL_image_shape)
+            print 'origin', origin_data.flatten()
+            print 'data', data.flatten()
+            image = Image.fromarray(data.reshape(shape))
+            # manully check the image and found that nothing wrong with the image storage.
+            # image.show()
 
-        # read and check whether the original image will be displayed
+            # after scale, elements are changed.
+            # self.assertTrue(
+            #     np.equal(origin_data.reshape(PIL_image_shape), data).all())
 
-        image_reader = self.reader.image(tag)
-        image_record = image_reader.record(0, 0)
-        data = image_record.data()
-        shape = image_record.shape()
+    def test_with_syntax(self):
+        with self.writer.mode("train") as writer:
+            scalar = writer.scalar("model/scalar/average")
+            for i in range(10):
+                scalar.add_record(i, float(i))
 
-        PIL_image_shape = (shape[0] * shape[1], shape[2])
-        data = np.array(data, dtype='uint8').reshape(PIL_image_shape)
-        print 'origin', origin_data.flatten()
-        print 'data', data.flatten()
-        image = Image.fromarray(data.reshape(shape))
-        # manully check the image and found that nothing wrong with the image storage.
-        # image.show()
+        self.reader = storage.LogReader(self.dir)
+        with self.reader.mode("train") as reader:
+            scalar = reader.scalar("model/scalar/average")
+            self.assertEqual(scalar.caption(), "train")
 
-        # after scale, elements are changed.
-        # self.assertTrue(
-        #     np.equal(origin_data.reshape(PIL_image_shape), data).all())
 
 
 if __name__ == '__main__':