Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.
(
fn,
*inputs,
grad=None,
repeats=10,
desc="",
verbose=True,
amp=False,
amp_dtype=torch.float16,
**kwinputs,
)
| 152 | |
| 153 | |
| 154 | def benchmark_all( |
| 155 | fn, |
| 156 | *inputs, |
| 157 | grad=None, |
| 158 | repeats=10, |
| 159 | desc="", |
| 160 | verbose=True, |
| 161 | amp=False, |
| 162 | amp_dtype=torch.float16, |
| 163 | **kwinputs, |
| 164 | ): |
| 165 | """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" |
| 166 | return ( |
| 167 | benchmark_forward( |
| 168 | fn, |
| 169 | *inputs, |
| 170 | repeats=repeats, |
| 171 | desc=desc, |
| 172 | verbose=verbose, |
| 173 | amp=amp, |
| 174 | amp_dtype=amp_dtype, |
| 175 | **kwinputs, |
| 176 | ), |
| 177 | benchmark_backward( |
| 178 | fn, |
| 179 | *inputs, |
| 180 | grad=grad, |
| 181 | repeats=repeats, |
| 182 | desc=desc, |
| 183 | verbose=verbose, |
| 184 | amp=amp, |
| 185 | amp_dtype=amp_dtype, |
| 186 | **kwinputs, |
| 187 | ), |
| 188 | benchmark_combined( |
| 189 | fn, |
| 190 | *inputs, |
| 191 | grad=grad, |
| 192 | repeats=repeats, |
| 193 | desc=desc, |
| 194 | verbose=verbose, |
| 195 | amp=amp, |
| 196 | amp_dtype=amp_dtype, |
| 197 | **kwinputs, |
| 198 | ), |
| 199 | ) |
| 200 | |
| 201 | |
| 202 | def pytorch_profiler( |
nothing calls this directly
no test coverage detected