send/recv

[daquexian]

用户接口的用法：

# rank 0:
x = flow.tensor([1, 2]).to('cuda')
flow.comm.send(x, 1)
# rank 1:
y = flow.comm.recv(0)
print(y)   # tensor([[1, 2]], device='cuda:1', dtype=oneflow.int64)

# rank 0:
x = flow.tensor([1, 2]).to('cuda')
flow.comm.send(x, 1, send_meta=False)
# rank 1:
y = flow.comm.recv(0, shape=(2,), dtype=flow.int64, device=flow.device('cuda'))
print(y)   # tensor([[1, 2]], device='cuda:1', dtype=oneflow.int64)

# rank 0:
x = flow.tensor([1, 2]).to('cuda')
flow.comm.send(x, 1)
# rank 1:
y = flow.zeros(2, dtype=flow.int64).to('cuda')
flow.comm.recv(0, out=y)
print(y)   # tensor([[1, 2]], device='cuda:1', dtype=oneflow.int64)

[github-actions]

Speed stats:

GPU Name: GeForce GTX 1080 

OneFlow resnet50 time: 127.6ms (= 6378.9ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 141.2ms (= 7062.0ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.11 (= 141.2ms / 127.6ms)

OneFlow resnet50 time: 74.4ms (= 3718.9ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 82.7ms (= 4134.9ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.11 (= 82.7ms / 74.4ms)

OneFlow resnet50 time: 47.0ms (= 2348.9ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 59.7ms (= 2985.4ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.27 (= 59.7ms / 47.0ms)

OneFlow resnet50 time: 39.4ms (= 1967.8ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 50.2ms (= 2508.0ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.27 (= 50.2ms / 39.4ms)

OneFlow resnet50 time: 40.1ms (= 2003.1ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 41.2ms (= 2059.4ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.03 (= 41.2ms / 40.1ms)

OneFlow resnet50 time: 139.5ms (= 6972.5ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 150.9ms (= 7545.7ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.08 (= 150.9ms / 139.5ms)

OneFlow resnet50 time: 87.0ms (= 4347.5ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 90.2ms (= 4509.7ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.04 (= 90.2ms / 87.0ms)

OneFlow resnet50 time: 62.6ms (= 3131.9ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 72.0ms (= 3599.7ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.15 (= 72.0ms / 62.6ms)

OneFlow resnet50 time: 58.8ms (= 2938.6ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 53.4ms (= 2672.1ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 0.91 (= 53.4ms / 58.8ms)

OneFlow resnet50 time: 51.3ms (= 2564.3ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 50.4ms (= 2519.4ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 0.98 (= 50.4ms / 51.3ms)

[github-actions]

Speed stats:

GPU Name: GeForce GTX 1080 

OneFlow resnet50 time: 127.9ms (= 6394.2ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 142.1ms (= 7107.1ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.11 (= 142.1ms / 127.9ms)

OneFlow resnet50 time: 74.5ms (= 3724.6ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 84.3ms (= 4214.7ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.13 (= 84.3ms / 74.5ms)

OneFlow resnet50 time: 47.3ms (= 2362.8ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 60.8ms (= 3039.8ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.29 (= 60.8ms / 47.3ms)

OneFlow resnet50 time: 41.3ms (= 2063.7ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 46.4ms (= 2321.1ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.12 (= 46.4ms / 41.3ms)

OneFlow resnet50 time: 38.6ms (= 1928.1ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 43.1ms (= 2154.5ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.12 (= 43.1ms / 38.6ms)

OneFlow resnet50 time: 139.5ms (= 6973.3ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 151.3ms (= 7563.0ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.08 (= 151.3ms / 139.5ms)

OneFlow resnet50 time: 88.2ms (= 4412.1ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 91.2ms (= 4561.9ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.03 (= 91.2ms / 88.2ms)

OneFlow resnet50 time: 62.7ms (= 3134.0ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 68.3ms (= 3415.1ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 1.09 (= 68.3ms / 62.7ms)

OneFlow resnet50 time: 62.3ms (= 3114.3ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 56.7ms (= 2833.8ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 0.91 (= 56.7ms / 62.3ms)

OneFlow resnet50 time: 60.9ms (= 3047.0ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
OneFlow GPU used (rank 0): 0 MiB
PyTorch resnet50 time: 47.5ms (= 2373.7ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
PyTorch GPU used (rank 0, estimated): 0 MiB
Relative speed: 0.78 (= 47.5ms / 60.9ms)