mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
bfeb45e46b
Follows #132604 Pull Request resolved: https://github.com/pytorch/pytorch/pull/132753 Approved by: https://github.com/Skylion007
28 lines
666 B
C++
28 lines
666 B
C++
#pragma once
|
|
|
|
#include <torch/csrc/jit/api/module.h>
|
|
|
|
namespace torch::jit {
|
|
|
|
struct TORCH_API LinearBNParameters {
|
|
at::Tensor linear_w;
|
|
at::Tensor linear_b;
|
|
at::Tensor bn_rm;
|
|
at::Tensor bn_rv;
|
|
double bn_eps = 0.0;
|
|
at::Tensor bn_w;
|
|
at::Tensor bn_b;
|
|
};
|
|
|
|
/**
|
|
* Given the current weight and bias tensors of a Linear module and parameters
|
|
* of the BatchNorm module we're folding with, compute the updated values
|
|
* for the weight and bias.
|
|
*
|
|
* The function is basically copied from torch/nn/utils/fusion.py
|
|
*/
|
|
TORCH_API std::tuple<at::Tensor, at::Tensor> computeUpdatedLinearWeightAndBias(
|
|
const LinearBNParameters& p);
|
|
|
|
} // namespace torch::jit
|