mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
956c22b1f9
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/62437 In this diff we add a broadcast fastpath for MulGradient and DivGradient ops, whose tests we update to exercise the new functionality. Test Plan: Added test cases to elementwise ops (which will exercise the new MulGradient / DivGradient broadcast fastpath functionality) that will be run by CI. It's worth noting there's still no code (outside of the new test cases) that takes the new code paths added -- the user must explicitly request allow_broadcast_fastpath=True, and nothing outside of the added tests currently does so. Differential Revision: D29938273 fbshipit-source-id: 281c1a109e38c25b9bf9ff8d832de60ac3c231a9
50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
|
|
|
|
|
|
|
|
|
|
import argparse
|
|
import numpy as np
|
|
|
|
from caffe2.python import core, workspace
|
|
|
|
|
|
def benchmark_mul_gradient(args):
|
|
workspace.FeedBlob("dC", np.random.rand(args.m, args.n).astype(np.float32))
|
|
workspace.FeedBlob("A", np.random.rand(args.m, args.n).astype(np.float32))
|
|
workspace.FeedBlob("B", np.random.rand(args.n).astype(np.float32))
|
|
|
|
net = core.Net("mynet")
|
|
net.MulGradient(
|
|
["dC", "A", "B"],
|
|
["dC" if args.inplace else "dA", "dB"],
|
|
broadcast=True,
|
|
axis=1,
|
|
allow_broadcast_fastpath=args.allow_broadcast_fastpath,
|
|
)
|
|
workspace.CreateNet(net)
|
|
|
|
workspace.BenchmarkNet(net.Name(), 1, args.iteration, True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
description="benchmark for MulGradient.")
|
|
parser.add_argument(
|
|
'-m', type=int, default=9508,
|
|
help="The number of rows of A")
|
|
parser.add_argument(
|
|
"-n", type=int, default=80,
|
|
help="The number of columns of A")
|
|
parser.add_argument(
|
|
'-i', "--iteration", type=int, default=100,
|
|
help="The number of iterations.")
|
|
parser.add_argument(
|
|
"--inplace",
|
|
action='store_true', help="Whether to perform the op in-place.")
|
|
parser.add_argument(
|
|
"--allow-broadcast-fastpath",
|
|
action='store_true', help="Whether the broadcast fastpath is enabled.")
|
|
args, extra_args = parser.parse_known_args()
|
|
core.GlobalInit(['python'] + extra_args)
|
|
benchmark_mul_gradient(args)
|