Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add oneflow.tensordot interface #7968

Merged
merged 155 commits into from
May 25, 2022
Merged

Add oneflow.tensordot interface #7968

merged 155 commits into from
May 25, 2022

Conversation

marigoold
Copy link
Contributor

@marigoold marigoold commented Apr 7, 2022

此PR完成了:

  • 增加了oneflow.tensordot的functional接口
  • 增加了oneflow.tensordot的单测和文档

算子实现的方法是在functor层,执行tensor permute -> tensor reshape -> functional::matmul -> tensor reshape的操作。

文档截图:
这里没写公式,因为我觉得那个非常难理解,而是加了一段代码的演示,告诉用户tensordot和哪一系列基本操作等价。

WX20220421-145607@2x

WX20220421-145620@2x

需要注意的地方:

  • 单测的test_tensordot_intdim中,dim=4的tensor会导致oom,所以dim设置为3
  • functor中的实现有很多中间变量,可能会造成内存问题(比如中间变量用了不必要的内存),需要麻烦各位reviewer重点看一下

@marigoold marigoold marked this pull request as ready for review April 7, 2022 07:11
@marigoold marigoold marked this pull request as draft April 7, 2022 08:19
@marigoold marigoold marked this pull request as ready for review April 18, 2022 01:41
@@ -659,6 +659,24 @@ def OneFlow_SiluGradOp : OneFlow_BaseOp<"silu_grad", [NoSideEffect, DeclareOpInt
let has_data_type_infer_fn = 1;
}

def OneFlow_TensorDotOp : OneFlow_BaseOp<"tensordot", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

不适用OpBuilder构建的op不需要更新这里

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

已修改

@marigoold marigoold requested a review from zhongshsh April 20, 2022 07:24


def tensordot(a, b, dims: Union[int, List[List[int]]] = 2):
if not isinstance(dims, (oneflow._oneflow_internal.Tensor, int, list, tuple)):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@marigoold marigoold requested review from oneflow-ci-bot and removed request for oneflow-ci-bot May 19, 2022 14:02
@github-actions
Copy link
Contributor

View latest API docs preview at: https://staging.oneflow.info/docs/Oneflow-Inc/oneflow/pr/7968/

@github-actions
Copy link
Contributor

CI failed when running job: cuda-module. PR label automerge has been removed

@github-actions
Copy link
Contributor

Speed stats:

@github-actions
Copy link
Contributor

View latest API docs preview at: https://staging.oneflow.info/docs/Oneflow-Inc/oneflow/pr/7968/

@github-actions
Copy link
Contributor

CI failed when running job: cpu-module. PR label automerge has been removed

@github-actions
Copy link
Contributor

CI failed when running job: cuda-module. PR label automerge has been removed

@github-actions
Copy link
Contributor

Speed stats:

@github-actions
Copy link
Contributor

Speed stats:

@github-actions
Copy link
Contributor

View latest API docs preview at: https://staging.oneflow.info/docs/Oneflow-Inc/oneflow/pr/7968/

@github-actions
Copy link
Contributor

Speed stats:

@github-actions
Copy link
Contributor

View latest API docs preview at: https://staging.oneflow.info/docs/Oneflow-Inc/oneflow/pr/7968/

@marigoold marigoold merged commit 50840cc into master May 25, 2022
@marigoold marigoold deleted the feat-add_tensordot_op branch May 25, 2022 07:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

9 participants