| 1347 | |
| 1348 | @lru_cache(maxsize=None) |
| 1349 | def _get_sync_bn_ops(device, dtype, eps_mode, ndim, channels): |
| 1350 | # fmt : off |
| 1351 | @subgraph("SyncBnStage0", dtype, device, 1) |
| 1352 | def syncbn_stage0(inputs, f, c): |
| 1353 | input = inputs[0] |
| 1354 | reduce_shape = c( |
| 1355 | (1, channels) + (1,) * (ndim - 2), dtype="int32", device=device |
| 1356 | ) |
| 1357 | input_shape = f(GetVarShape(), input) |
| 1358 | input_elems = f(Reduce(mode="product", axis=0), input_shape) |
| 1359 | reduce_elems = f(Reduce(mode="product", axis=0), reduce_shape) |
| 1360 | reduce_size = f("//", input_elems, reduce_elems) |
| 1361 | channel_x1s = f(Reduce(mode="sum"), input, reduce_shape) |
| 1362 | channel_x2s = f(Reduce(mode="sum_sqr"), input, reduce_shape) |
| 1363 | reduce_size_f = f(TypeCvt(dtype=dtype), reduce_size) |
| 1364 | return ( |
| 1365 | (reduce_shape, reduce_size_f, channel_x1s, channel_x2s), |
| 1366 | (False, False, True, True), |
| 1367 | ) |
| 1368 | |
| 1369 | @subgraph("SyncBnStage1", dtype, device, 7) |
| 1370 | def syncbn_stage1(inputs, f, c): |
| 1371 | input, reduce_size, channel_x1s, channel_x2s, eps = inputs[0:5] |
| 1372 | weight, bias = inputs[5:7] |
| 1373 | channel_mean = f("/", channel_x1s, reduce_size) |
| 1374 | channel_var = f( |
| 1375 | "+", |
| 1376 | f( |
| 1377 | "/", |
| 1378 | f("**", channel_x1s, c(2)), |
| 1379 | f("-", f("*", reduce_size, reduce_size)), |
| 1380 | ), |
| 1381 | f("/", channel_x2s, reduce_size), |
| 1382 | ) |
| 1383 | invsqrt_channel_var = f("**", f(eps_mode, channel_var, eps), c(-0.5)) |
| 1384 | inv_var_wt = f("*", invsqrt_channel_var, weight) |
| 1385 | neg_channel_mean = f("-", channel_mean) |
| 1386 | outvar = f( |
| 1387 | "fma3", |
| 1388 | input, |
| 1389 | inv_var_wt, |
| 1390 | f("+", f("*", neg_channel_mean, inv_var_wt), bias), |
| 1391 | ) |
| 1392 | return (outvar, channel_mean, channel_var), (True, True, True) |
| 1393 | |
| 1394 | @subgraph("SyncBnStage1Inference", dtype, device, 6) |
| 1395 | def syncbn_stage1_inference(inputs, f, c): |
| 1396 | input, channel_mean, channel_var, eps = inputs[0:4] |
| 1397 | weight, bias = inputs[4:6] |
| 1398 | invsqrt_channel_var = f("**", f(eps_mode, channel_var, eps), c(-0.5)) |
| 1399 | inv_var_wt = f("*", invsqrt_channel_var, weight) |
| 1400 | neg_channel_mean = f("-", channel_mean) |
| 1401 | outvar = f( |
| 1402 | "+", |
| 1403 | f("*", input, inv_var_wt), |
| 1404 | f("+", f("*", neg_channel_mean, inv_var_wt), bias), |
| 1405 | ) |
| 1406 | return (outvar,), (True,) |