1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
use crate::paramstore::{ParamStore, ParamStoreError};
use nove_tensor::{Device, Tensor, TensorError};
use thiserror::Error;
pub mod layer;
pub mod paramstore;
#[derive(Error, Debug)]
pub enum ModelError {
#[error("IO error: {0}")]
IoError(#[from] std::io::Error),
#[error("Tensor error: {0}")]
TensorError(#[from] TensorError),
#[error("ParamStore error: {0}")]
ParamStoreError(#[from] ParamStoreError),
#[error("Other error: {0}")]
OtherError(String),
}
#[derive(Debug, Clone)]
pub struct Parameter(pub String, pub Tensor);
pub trait Model {
type ParamStore: ParamStore;
type Input;
type Output;
/// Get the direct parameter stores of the model.
///
/// # Returns
/// * `Ok(Vec<Self::ParamStore>)` - The parameter stores of the model.
/// * `Err(ModelError)` - The error when getting the parameter stores.
fn param_stores(&self) -> Result<Vec<Self::ParamStore>, ModelError>;
/// Get all parameters of the model, including those in the submodules.
///
/// # Returns
/// * `Ok(Vec<Parameter>)` - All parameters of the model.
/// * `Err(ModelError)` - The error when getting the parameters.
fn parameters(&self) -> Result<Vec<Parameter>, ModelError> {
let mut all_params = Vec::new();
fn collect_params<S: ParamStore>(
store: &S,
all_params: &mut Vec<Parameter>,
) -> Result<(), ParamStoreError> {
// Collect parameters from the current store.
all_params.extend(store.parameters()?);
// Recursively collect parameters from submodules.
for submodule in store.modules()? {
collect_params(&submodule, all_params)?;
}
Ok(())
}
for store in self.param_stores()? {
collect_params(&store, &mut all_params)?;
}
Ok(all_params)
}
/// Perform a forward pass of the model.
///
/// # Arguments
/// * `input` - The input data for the forward pass.
///
/// # Returns
/// * `Ok(Self::Output)` - The output data of the forward pass.
/// * `Err(ModelError)` - The error when performing the forward pass.
fn forward(&mut self, input: Self::Input) -> Result<Self::Output, ModelError>;
/// Save the model to the specified folder.
///
/// # Notes
/// * This method would call the `save` method of each parameter store. So the specific storage
/// method depends on the ParamStore chosen for the model.
///
/// # Arguments
/// * `folder_path` - The path of the folder to save the model to.
///
/// # Returns
/// * `Ok(())` - If the model is successfully saved to the folder.
/// * `Err(ModelError)` - The error when saving the model to the folder.
fn save(&self, folder_path: &str) -> Result<(), ModelError> {
// Create the folder if it does not exist.
std::fs::create_dir_all(folder_path)?;
let param_stores = self.param_stores()?;
for param_store in param_stores {
param_store.save(folder_path)?;
}
Ok(())
}
/// Load the model from the specified folder.
///
/// # Notes
/// * This method would call the `load` method of each parameter store.
/// * The ParamStore type must be same as the one used for saving.
///
/// # Arguments
/// * `folder_path` - The path of the folder to load the model from.
/// * `devices` - The devices to load the model to.
/// * If the number of devices is 1, the model would be loaded to all parameter stores to this device.
/// * If the number of devices is equal to the number of parameter stores, each parameter store
/// would be loaded to the corresponding device.
/// * If the number of devices is not equal to 1 or the number of parameter stores, the method
/// would return an error.
/// * `process_fn` - The function to process each parameter store(including submodules).
///
/// # Returns
/// * `Ok(())` - If the model is successfully loaded from the folder.
/// * `Err(ModelError)` - The error when loading the model from the folder.
fn load<F>(
&mut self,
folder_path: &str,
devices: &[Device],
mut process_fn: F,
) -> Result<(), ModelError>
where
F: FnMut(&str, &Self::ParamStore) -> Result<(), ParamStoreError>,
{
if devices.is_empty() {
return Err(ModelError::OtherError(
"The devices list is empty.".to_string(),
));
}
let param_stores = self.param_stores()?;
// Get the number of parameter stores.
let num_param_stores = param_stores.len();
if num_param_stores == 0 {
return Err(ModelError::OtherError(
"The model has no parameter store.".to_string(),
));
}
if devices.len() != 1 && devices.len() != num_param_stores {
return Err(ModelError::OtherError(format!(
"The number of devices({}) is not equal to the number of parameter stores({}) or 1.",
devices.len(),
num_param_stores
)));
}
match devices.len() {
1 => {
for param_store in param_stores.iter() {
param_store.load(folder_path, &devices[0], &mut process_fn)?;
}
}
_ => {
for (i, param_store) in param_stores.iter().enumerate() {
param_store.load(folder_path, &devices[i], &mut process_fn)?;
}
}
}
Ok(())
}
/// Get the summary of the model.
///
/// # Returns
/// * `Ok(String)` - The summary of the model.
/// * `Err(ModelError)` - The error when getting the model summary.
fn summary(&self) -> Result<String, ModelError> {
let mut summary = format!("{}(\n", std::any::type_name::<Self>());
let param_stores = self.param_stores()?;
for param_store in param_stores {
let param_store_summary = format!("{}", param_store);
for line in param_store_summary.lines() {
summary.push_str(&format!(" {}\n", line));
}
}
summary.push_str(")\n");
Ok(summary)
}
}