[JIT] AutoBatching - IR transformation(basic operators) by ChunliF · Pull Request #9198 · pytorch/pytorch

ChunliF · 2018-07-06T00:08:39Z

Use decorator torch.jit.batch to implement auto-batching (call to_batch pass to do IR tranformation).

to_batch pass: "to_batch.h/cpp" in csrc/jit/passess to transform a graph to a new batched graph.
Write several basic operators for BatchTensor (add, mul, sigmoid, tanh, mm, matmul, select).
Register the operators in a lookup table <std::string, std::shared_ptr<Graph>>. (use the Graph to replace the original node in IR graph)

Move BatchTensor in python from torch.BatchTensor to torch.jit.BatchTensor

facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

torch/csrc/jit/passes/to_batch.cpp

jamesr66a

Looking good!

torch/jit/__init__.py

torch/csrc/jit/passes/to_batch.cpp

torch/jit/__init__.py

facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

jamesr66a

Looks good!

zdevito

This looks pretty good! I mostly have code organizational comments and a few nits.

torch/csrc/jit/passes/to_batch.cpp

+}
+
+void to_batch_graph(std::shared_ptr<Graph>& graph, int64_t batch_size){
+  // batch_size: not used yet, will be used to deal with scalarType


torch/csrc/jit/passes/to_batch.cpp

+namespace torch { namespace jit {
+
+// map from batchTensor to {data, mask, dims}
+static std::unordered_map<Value*, std::vector<Value*>> batch_map;


torch/csrc/jit/passes/to_batch.cpp

+          auto new_input = batch_map.at(input);
+          new_inputs.insert(new_inputs.end(), new_input.begin(), new_input.end());
+        }
+        else{


torch/csrc/jit/passes/to_batch.cpp

+        batch_map[output] = std::vector<Value*>(outputs.begin() + i * 3, outputs.begin() + i * 3 + 3);
+      }
+    }
+    // control flow: not supported yet, will be added further


torch/csrc/jit/script/init.cpp

        // There is a thin wrapper on top of this method in the C++ version of
        // ScriptModule.
        return runMethodFromPython(self.get_method("forward"), args);
+      })


torch/jit/__init__.py

+def batch(batch_size=1, optimize=True, _frames_up=0):
+    def decorator(fn):
+        mod = script(fn, optimize, _frames_up)
+        res_graph = mod.to_batch(batch_size)


facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@ChunliF has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@ChunliF is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

Summary: Use decorator `torch.jit.batch` to implement auto-batching (call `to_batch` pass to do IR tranformation). - `to_batch` pass: "to_batch.h/cpp" in csrc/jit/passess to transform a graph to a new batched graph. - Write several basic operators for BatchTensor (add, mul, sigmoid, tanh, mm, matmul, select). - Register the operators in a lookup table `<std::string, std::shared_ptr<Graph>>`. (use the Graph to replace the original node in IR graph) Move BatchTensor in python from torch.BatchTensor to torch.jit.BatchTensor Pull Request resolved: pytorch#9198 Reviewed By: zdevito Differential Revision: D8744466 Pulled By: ChunliF fbshipit-source-id: 9ea56a30f55cb870f13a2069a47cc635419763ff

ChunliF added 6 commits June 26, 2018 16:14

add BatchTensor class

efa2475

change batchtensor.dims to tensor

d8a1ea8

add comment, fix lint

3fb9455

fix merge

3827725

fix merge conflict

90030c0

ir transformation - basic operators

1576966

ChunliF requested review from apaszke, colesbury, ezyang, gchanan, soumith and zdevito as code owners July 6, 2018 00:08

facebook-github-bot reviewed Jul 6, 2018

View reviewed changes

jamesr66a reviewed Jul 6, 2018

View reviewed changes

torch/csrc/jit/passes/to_batch.cpp Outdated

This comment was marked as off-topic.

Sign in to view

jamesr66a reviewed Jul 6, 2018

View reviewed changes

torch/csrc/jit/passes/to_batch.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

jamesr66a reviewed Jul 6, 2018

View reviewed changes

torch/csrc/jit/passes/to_batch.cpp Outdated

This comment was marked as off-topic.

Sign in to view

jamesr66a reviewed Jul 6, 2018

View reviewed changes

torch/csrc/jit/passes/to_batch.cpp Outdated

This comment was marked as off-topic.

Sign in to view

jamesr66a reviewed Jul 6, 2018

View reviewed changes

ChunliF added 2 commits July 6, 2018 11:34

fix merge

ff24d4a

fix bugs

fb923c6

ChunliF force-pushed the batchir branch from d36c7b1 to fb923c6 Compare July 6, 2018 18:36

facebook-github-bot reviewed Jul 6, 2018

View reviewed changes

ChunliF added 2 commits July 6, 2018 15:52

move batch_operators to a new file

bb858d6

to_batch pass creates new graph, add lstm test case

fcedaf2

facebook-github-bot reviewed Jul 9, 2018

View reviewed changes

jamesr66a approved these changes Jul 9, 2018

View reviewed changes

move BatchTensor to torch.jit, try to fix lint

92923ca

ChunliF mentioned this pull request Jul 10, 2018

[JIT] Add BatchTensor class #8922

Closed

fmassa mentioned this pull request Jul 10, 2018

concat tensor #9298

Closed

zdevito reviewed Jul 11, 2018

View reviewed changes

ChunliF and others added 6 commits July 11, 2018 13:27

create ToBatch class

03f936b

fix merge conflict

1726ec7

Merge branch 'master' into batchir

487ab60

fix bug

9fec40e

fix merge conflict

871aea0

Merge branch 'batchir' of github.com:ChunliF/pytorch into batchir

9ba178c

facebook-github-bot reviewed Jul 11, 2018

View reviewed changes

ChunliF added 3 commits July 11, 2018 14:27

attempt to fix build error

f7ba2cb

remove to_batch method from ScriptModule

ca64e14

fix lint

ab6f16c

facebook-github-bot reviewed Jul 11, 2018

View reviewed changes

fix bug

e2d6ade

facebook-github-bot reviewed Jul 11, 2018

View reviewed changes

facebook-github-bot closed this in a487b08 Jul 12, 2018

ChunliF deleted the batchir branch July 12, 2018 17:42

ezyang added the merged label Jun 26, 2019

Conversation

ChunliF commented Jul 6, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

jamesr66a left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

jamesr66a left a comment

Choose a reason for hiding this comment

Uh oh!

zdevito left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

ChunliF commented Jul 6, 2018 •

edited

Loading