Merge pull request #47 from ChunweiYan/feature/refactor_python_sdk

server/visualdl/lib.py

@@ -16,7 +16,7 @@ def get_modes(storage):
 def get_scalar_tags(storage, mode):
     result = {}
     for mode in storage.modes():
-        reader = storage.as_mode(mode)
+        with storage.mode(mode) as reader:
             tags = reader.tags('scalar')
             if tags:
                 result[mode] = {}
@@ -29,7 +29,7 @@ def get_scalar_tags(storage, mode):
 
 
 def get_scalar(storage, mode, tag):
-    reader = storage.as_mode(mode)
+    with storage.mode(mode) as reader:
         scalar = reader.scalar(tag)
 
         records = scalar.records()
@@ -44,7 +44,7 @@ def get_image_tags(storage):
     result = {}
 
     for mode in storage.modes():
-        reader = storage.as_mode(mode)
+        with storage.mode(mode) as reader:
             tags = reader.tags('image')
             if tags:
                 result[mode] = {}
@@ -70,7 +70,7 @@ def get_image_tag_steps(storage, mode, tag):
         tag = tag[:tag.rfind('/')]
         sample_index = int(res.groups()[0])
 
-    reader = storage.as_mode(mode)
+    with storage.mode(mode) as reader:
         image = reader.image(tag)
         res = []
 
@@ -95,7 +95,7 @@ def get_image_tag_steps(storage, mode, tag):
 
 
 def get_invididual_image(storage, mode, tag, step_index):
-    reader = storage.as_mode(mode)
+    with storage.mode(mode) as reader:
         res = re.search(r".*/([0-9]+$)", tag)
         # remove suffix '/x'
         if res:
@@ -114,7 +114,7 @@ def get_invididual_image(storage, mode, tag, step_index):
 
 
 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,7 +6,7 @@ import numpy as np
 
 
 def add_scalar(writer, mode, tag, num_steps, skip):
-    my_writer = writer.as_mode(mode)
+    with writer.mode(mode) as my_writer:
         scalar = my_writer.scalar(tag)
         for i in range(num_steps):
             if i % skip == 0:
@@ -20,7 +20,7 @@ def add_image(writer,
               num_passes,
               step_cycle,
               shape=[50, 50, 3]):
-    writer_ = writer.as_mode(mode)
+    with writer.mode(mode) as writer_:
         image_writer = writer_.image(tag, num_samples, step_cycle)
 
         for pass_ in xrange(num_passes):

visualdl/logic/pybind.cc

@@ -12,48 +12,48 @@ 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) { \
+  .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) { \
+  .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,8 +22,9 @@ 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.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()
@@ -49,8 +50,9 @@ 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.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)
 
@@ -59,7 +61,7 @@ class StorageTest(unittest.TestCase):
             data = image_record.data()
             self.assertEqual(len(data), np.prod(shape))
 
-        image_tags = self.reader.tags("image")
+            image_tags = reader.tags("image")
             self.assertTrue(image_tags)
             self.assertEqual(len(image_tags), 1)
 
@@ -75,7 +77,8 @@ 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()
@@ -84,7 +87,7 @@ class StorageTest(unittest.TestCase):
 
             # read and check whether the original image will be displayed
 
-        image_reader = self.reader.image(tag)
+            image_reader = reader.image(tag)
             image_record = image_reader.record(0, 0)
             data = image_record.data()
             shape = image_record.shape()
@@ -101,6 +104,18 @@ class StorageTest(unittest.TestCase):
             # self.assertTrue(
             #     np.equal(origin_data.reshape(PIL_image_shape), data).all())
 
+    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))
+
+        self.reader = storage.LogReader(self.dir)
+        with self.reader.mode("train") as reader:
+            scalar = reader.scalar("model/scalar/average")
+            self.assertEqual(scalar.caption(), "train")
+
+
 
 if __name__ == '__main__':
     unittest.main()