Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
86 lines
2.6 KiB
C++
86 lines
2.6 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
|
|
#include <cstdint>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
|
|
namespace engine::assets {
|
|
|
|
template <typename BatchNorm1dWeights>
|
|
BatchNorm1dWeights load_batch_norm1d_metadata(
|
|
const TensorSource & source,
|
|
const std::string & prefix) {
|
|
BatchNorm1dWeights bn;
|
|
const auto weight = source.require_metadata(prefix + ".weight");
|
|
if (weight.shape.size() != 1) {
|
|
throw std::runtime_error("batch norm weight must be 1D at " + prefix);
|
|
}
|
|
bn.channels = weight.shape[0];
|
|
bn.weight = source.require_f32(prefix + ".weight", {bn.channels});
|
|
bn.bias = source.require_f32(prefix + ".bias", {bn.channels});
|
|
bn.running_mean = source.require_f32(prefix + ".running_mean", {bn.channels});
|
|
bn.running_var = source.require_f32(prefix + ".running_var", {bn.channels});
|
|
return bn;
|
|
}
|
|
|
|
template <typename Conv1dWeights>
|
|
Conv1dWeights load_conv1d_metadata(
|
|
const TensorSource & source,
|
|
const std::string & prefix,
|
|
int64_t stride,
|
|
int64_t dilation,
|
|
int64_t padding,
|
|
int64_t groups,
|
|
bool use_bias) {
|
|
Conv1dWeights conv;
|
|
const auto weight = source.require_metadata(prefix + ".weight");
|
|
if (weight.shape.size() != 3) {
|
|
throw std::runtime_error("conv weight must be 3D at " + prefix);
|
|
}
|
|
conv.weight_name = prefix + ".weight";
|
|
conv.weight_source_shape = weight.shape;
|
|
conv.out_channels = weight.shape[0];
|
|
conv.in_channels = weight.shape[1] * groups;
|
|
conv.kernel = weight.shape[2];
|
|
conv.stride = stride;
|
|
conv.dilation = dilation;
|
|
conv.padding = padding;
|
|
conv.groups = groups;
|
|
conv.use_bias = use_bias;
|
|
if (use_bias) {
|
|
(void) source.require_metadata(prefix + ".bias");
|
|
conv.bias_name = prefix + ".bias";
|
|
}
|
|
return conv;
|
|
}
|
|
|
|
template <typename Conv1dWeights>
|
|
Conv1dWeights load_linear_as_conv1d_metadata(
|
|
const TensorSource & source,
|
|
const std::string & prefix,
|
|
bool use_bias) {
|
|
Conv1dWeights conv;
|
|
const auto weight = source.require_metadata(prefix + ".weight");
|
|
if (weight.shape.size() != 2) {
|
|
throw std::runtime_error("linear weight must be 2D at " + prefix);
|
|
}
|
|
conv.weight_name = prefix + ".weight";
|
|
conv.weight_source_shape = weight.shape;
|
|
conv.out_channels = weight.shape[0];
|
|
conv.in_channels = weight.shape[1];
|
|
conv.kernel = 1;
|
|
conv.stride = 1;
|
|
conv.dilation = 1;
|
|
conv.padding = 0;
|
|
conv.groups = 1;
|
|
conv.use_bias = use_bias;
|
|
if (use_bias) {
|
|
(void) source.require_metadata(prefix + ".bias");
|
|
conv.bias_name = prefix + ".bias";
|
|
}
|
|
return conv;
|
|
}
|
|
|
|
} // namespace engine::assets
|