dev_scaled_dot_product_attention_math

[ccssu]

Tensor attn_mask=None, Float dropout_p=0.0, Bool is_causal=False, Tensor dropout_mask=None

Todo

测试

• attention(query, key, value)
• attn_mask
• dropout_p , dropout_mask
• is_causal
• 添加测试函数

文档

• 使用指南

[github-actions]

View latest API docs preview at: https://oneflow-staging.oss-cn-beijing.aliyuncs.com/docs/Oneflow-Inc/oneflow/pr/10333/

[github-actions]

Speed stats:

GPU Name: NVIDIA GeForce RTX 3080 Ti 

❌ OneFlow resnet50 time: 43.7ms (= 4371.9ms / 100, input_shape=[16, 3, 224, 224])
PyTorch resnet50 time: 61.6ms (= 6155.2ms / 100, input_shape=[16, 3, 224, 224])
✔️ Relative speed: 1.41 (= 61.6ms / 43.7ms)

OneFlow resnet50 time: 26.1ms (= 2607.4ms / 100, input_shape=[8, 3, 224, 224])
PyTorch resnet50 time: 37.8ms (= 3783.7ms / 100, input_shape=[8, 3, 224, 224])
✔️ Relative speed: 1.45 (= 37.8ms / 26.1ms)

OneFlow resnet50 time: 19.4ms (= 3874.7ms / 200, input_shape=[4, 3, 224, 224])
PyTorch resnet50 time: 35.5ms (= 7101.9ms / 200, input_shape=[4, 3, 224, 224])
✔️ Relative speed: 1.83 (= 35.5ms / 19.4ms)

OneFlow resnet50 time: 17.7ms (= 3535.4ms / 200, input_shape=[2, 3, 224, 224])
PyTorch resnet50 time: 33.2ms (= 6631.6ms / 200, input_shape=[2, 3, 224, 224])
✔️ Relative speed: 1.88 (= 33.2ms / 17.7ms)

OneFlow resnet50 time: 17.3ms (= 3452.3ms / 200, input_shape=[1, 3, 224, 224])
PyTorch resnet50 time: 29.6ms (= 5911.7ms / 200, input_shape=[1, 3, 224, 224])
✔️ Relative speed: 1.71 (= 29.6ms / 17.3ms)

OneFlow swin dataloader time: 0.202s (= 40.392s / 200, num_workers=1)
PyTorch swin dataloader time: 0.130s (= 25.951s / 200, num_workers=1)
Relative speed: 0.642 (= 0.130s / 0.202s)

OneFlow swin dataloader time: 0.056s (= 11.192s / 200, num_workers=4)
PyTorch swin dataloader time: 0.033s (= 6.500s / 200, num_workers=4)
Relative speed: 0.581 (= 0.033s / 0.056s)

OneFlow swin dataloader time: 0.031s (= 6.293s / 200, num_workers=8)
PyTorch swin dataloader time: 0.016s (= 3.288s / 200, num_workers=8)
Relative speed: 0.522 (= 0.016s / 0.031s)

❌ OneFlow resnet50 time: 47.6ms (= 4756.2ms / 100, input_shape=[16, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 64.0ms (= 6400.7ms / 100, input_shape=[16, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.35 (= 64.0ms / 47.6ms)

OneFlow resnet50 time: 31.2ms (= 3124.3ms / 100, input_shape=[8, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 45.9ms (= 4591.6ms / 100, input_shape=[8, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.47 (= 45.9ms / 31.2ms)

OneFlow resnet50 time: 24.1ms (= 4820.5ms / 200, input_shape=[4, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 41.5ms (= 8297.7ms / 200, input_shape=[4, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.72 (= 41.5ms / 24.1ms)

OneFlow resnet50 time: 22.1ms (= 4413.9ms / 200, input_shape=[2, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 36.5ms (= 7307.6ms / 200, input_shape=[2, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.66 (= 36.5ms / 22.1ms)

OneFlow resnet50 time: 21.1ms (= 4214.1ms / 200, input_shape=[1, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 34.3ms (= 6860.4ms / 200, input_shape=[1, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.63 (= 34.3ms / 21.1ms)