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
71 lines
2.2 KiB
C++
71 lines
2.2 KiB
C++
#pragma once
|
|
|
|
#include "engine/models/niagara_asr/assets.h"
|
|
#include "engine/framework/core/backend_weight_store.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
#include "engine/framework/modules/conv_modules.h"
|
|
#include "engine/framework/modules/linear_module.h"
|
|
#include "engine/framework/modules/norm_modules.h"
|
|
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
namespace engine::models::niagara_asr {
|
|
|
|
struct NiagaraSubsamplingWeights {
|
|
engine::modules::BatchNorm2dEvalWeights bn0;
|
|
engine::modules::BatchNorm2dEvalWeights bn1;
|
|
engine::modules::Conv2dWeights input_projection;
|
|
engine::modules::Conv2dWeights conv1;
|
|
engine::modules::Conv2dWeights conv2;
|
|
engine::modules::LinearWeights dense;
|
|
};
|
|
|
|
struct NiagaraSelfAttentionWeights {
|
|
engine::modules::NormWeights norm;
|
|
engine::modules::LinearWeights query;
|
|
engine::modules::LinearWeights key;
|
|
engine::modules::LinearWeights value;
|
|
engine::modules::LinearWeights global_attention;
|
|
engine::modules::LinearWeights out;
|
|
engine::core::TensorValue relative_position;
|
|
engine::core::TensorValue u_bias;
|
|
engine::core::TensorValue v_bias;
|
|
};
|
|
|
|
struct NiagaraFeedForwardWeights {
|
|
engine::modules::NormWeights norm;
|
|
engine::modules::LinearWeights dense1;
|
|
engine::modules::LinearWeights dense2;
|
|
engine::core::TensorValue residual_scale;
|
|
};
|
|
|
|
struct NiagaraStateSpaceWeights {
|
|
engine::modules::NormWeights norm;
|
|
engine::modules::LinearWeights dense1;
|
|
engine::modules::LinearWeights dense2;
|
|
engine::core::TensorValue conv_kernel;
|
|
};
|
|
|
|
struct NiagaraLayerWeights {
|
|
NiagaraFeedForwardWeights ffn;
|
|
NiagaraSelfAttentionWeights self_attention;
|
|
NiagaraStateSpaceWeights state_space;
|
|
NiagaraFeedForwardWeights ffn1;
|
|
engine::modules::NormWeights out_norm;
|
|
};
|
|
|
|
struct NiagaraWeights {
|
|
std::shared_ptr<engine::core::BackendWeightStore> store;
|
|
NiagaraSubsamplingWeights subsampling;
|
|
std::vector<NiagaraLayerWeights> layers;
|
|
engine::modules::LinearWeights ctc;
|
|
};
|
|
|
|
std::shared_ptr<const NiagaraWeights> load_niagara_weights(
|
|
const NiagaraAsrAssets & assets,
|
|
engine::core::ExecutionContext & execution,
|
|
engine::assets::TensorStorageType storage_type,
|
|
size_t weight_context_bytes);
|
|
|
|
} // namespace engine::models::niagara_asr
|