1// This file contains functionality for training models using GGML.
2// It is not strictly needed vs. just vanilla GGML but it provides a more high-level interface for common needs such as datasets.
3// At the bottom of this file especially there are relatively high-level functions that are suitable use or adaptation in user code.
4//
5// Module maintainer: Johannes Gäßler (@JohannesGaessler, johannesg@5d6.de)
6
7#pragma once
8
9#include "ggml.h"
10#include "ggml-backend.h"
11
12#include <stdint.h>
13
14#ifdef __cplusplus
15extern "C" {
16#endif
17
18 struct ggml_opt_dataset;
19 struct ggml_opt_context;
20 struct ggml_opt_result;
21
22 typedef struct ggml_opt_dataset * ggml_opt_dataset_t;
23 typedef struct ggml_opt_context * ggml_opt_context_t;
24 typedef struct ggml_opt_result * ggml_opt_result_t;
25
26 // ====== Loss ======
27
28 // built-in loss types, i.e. the built-in quantities minimized by the optimizer
29 // custom loss types can be defined via mean or sum which simply reduce the outputs for all datapoints to a single value
30 enum ggml_opt_loss_type {
31 GGML_OPT_LOSS_TYPE_MEAN,
32 GGML_OPT_LOSS_TYPE_SUM,
33 GGML_OPT_LOSS_TYPE_CROSS_ENTROPY,
34 GGML_OPT_LOSS_TYPE_MEAN_SQUARED_ERROR,
35 };
36
37 // ====== Dataset ======
38
39 GGML_API ggml_opt_dataset_t ggml_opt_dataset_init(
40 enum ggml_type type_data, // the type for the internal data tensor
41 enum ggml_type type_label, // the type for the internal labels tensor
42 int64_t ne_datapoint, // number of elements per datapoint
43 int64_t ne_label, // number of elements per label
44 int64_t ndata, // total number of datapoints/labels
45 int64_t ndata_shard); // number of datapoints/labels per shard (unit at which the dataset is shuffled/copied)
46 GGML_API void ggml_opt_dataset_free(ggml_opt_dataset_t dataset);
47
48 // get underlying tensors that store the data
49 GGML_API int64_t ggml_opt_dataset_ndata (ggml_opt_dataset_t dataset);
50 GGML_API struct ggml_tensor * ggml_opt_dataset_data (ggml_opt_dataset_t dataset); // shape = [ne_datapoint, ndata]
51 GGML_API struct ggml_tensor * ggml_opt_dataset_labels(ggml_opt_dataset_t dataset); // shape = [nd_label, ndata]
52
53 // shuffle idata first datapoints from dataset with RNG from opt_ctx, shuffle all datapoints if idata is negative
54 GGML_API void ggml_opt_dataset_shuffle(ggml_opt_context_t opt_ctx, ggml_opt_dataset_t dataset, int64_t idata);
55
56 // get batch at position ibatch from dataset and copy the data to data_batch and labels_batch
57 GGML_API void ggml_opt_dataset_get_batch(
58 ggml_opt_dataset_t dataset,
59 struct ggml_tensor * data_batch, // shape = [ne_datapoint, ndata_batch]
60 struct ggml_tensor * labels_batch, // shape = [ne_label, ndata_batch]
61 int64_t ibatch);
62 GGML_API void ggml_opt_dataset_get_batch_host(
63 ggml_opt_dataset_t dataset,
64 void * data_batch,
65 size_t nb_data_batch,
66 void * labels_batch,
67 int64_t ibatch);
68
69 // ====== Model / Context ======
70
71 enum ggml_opt_build_type {
72 GGML_OPT_BUILD_TYPE_FORWARD = 10,
73 GGML_OPT_BUILD_TYPE_GRAD = 20,
74 GGML_OPT_BUILD_TYPE_OPT = 30,
75 };
76
77 enum ggml_opt_optimizer_type {
78 GGML_OPT_OPTIMIZER_TYPE_ADAMW,
79 GGML_OPT_OPTIMIZER_TYPE_SGD,
80
81 GGML_OPT_OPTIMIZER_TYPE_COUNT
82 };
83
84 // parameters that control which optimizer is used and how said optimizer tries to find the minimal loss
85 struct ggml_opt_optimizer_params {
86 struct {
87 float alpha; // learning rate
88 float beta1; // first AdamW momentum
89 float beta2; // second AdamW momentum
90 float eps; // epsilon for numerical stability
91 float wd; // weight decay - 0.0f to disable
92 } adamw;
93 struct {
94 float alpha; // learning rate
95 float wd; // weight decay
96 } sgd;
97 };
98
99 // callback to calculate optimizer parameters prior to a backward pass
100 // userdata can be used to pass arbitrary data
101 typedef struct ggml_opt_optimizer_params (*ggml_opt_get_optimizer_params)(void * userdata);
102
103 // returns the default optimizer params (constant, hard-coded values)
104 // userdata is not used
105 GGML_API struct ggml_opt_optimizer_params ggml_opt_get_default_optimizer_params(void * userdata);
106
107 // casts userdata to ggml_opt_optimizer_params and returns it
108 GGML_API struct ggml_opt_optimizer_params ggml_opt_get_constant_optimizer_params(void * userdata);
109
110 // parameters for initializing a new optimization context
111 struct ggml_opt_params {
112 ggml_backend_sched_t backend_sched; // defines which backends are used to construct the compute graphs
113
114 // by default the forward graph needs to be reconstructed for each eval
115 // if ctx_compute, inputs, and outputs are set the graphs are instead allocated statically
116 struct ggml_context * ctx_compute;
117 struct ggml_tensor * inputs;
118 struct ggml_tensor * outputs;
119
120 enum ggml_opt_loss_type loss_type;
121 enum ggml_opt_build_type build_type;
122
123 int32_t opt_period; // after how many gradient accumulation steps an optimizer step should be done
124
125 ggml_opt_get_optimizer_params get_opt_pars; // callback for calculating optimizer parameters
126 void * get_opt_pars_ud; // userdata for calculating optimizer parameters
127
128 // only GGML_OPT_OPTIMIZER_TYPE_ADAMW needs m, v momenta per parameter tensor
129 enum ggml_opt_optimizer_type optimizer;
130 };
131
132 // get parameters for an optimization context with defaults set where possible
133 // parameters for which no sensible defaults exist are supplied as arguments to this function
134 GGML_API struct ggml_opt_params ggml_opt_default_params(
135 ggml_backend_sched_t backend_sched,
136 enum ggml_opt_loss_type loss_type);
137
138 GGML_API ggml_opt_context_t ggml_opt_init(struct ggml_opt_params params);
139 GGML_API void ggml_opt_free(ggml_opt_context_t opt_ctx);
140
141 // set gradients to zero, initilize loss, and optionally reset the optimizer
142 GGML_API void ggml_opt_reset(ggml_opt_context_t opt_ctx, bool optimizer);
143
144 GGML_API bool ggml_opt_static_graphs(ggml_opt_context_t opt_ctx); // whether the graphs are allocated_statically
145
146 // get underlying tensors that store data
147 // if not using static graphs these pointers become invalid with the next call to ggml_opt_alloc
148 GGML_API struct ggml_tensor * ggml_opt_inputs( ggml_opt_context_t opt_ctx); // forward graph input tensor
149 GGML_API struct ggml_tensor * ggml_opt_outputs( ggml_opt_context_t opt_ctx); // forward graph output tensor
150 GGML_API struct ggml_tensor * ggml_opt_labels( ggml_opt_context_t opt_ctx); // labels to compare outputs against
151 GGML_API struct ggml_tensor * ggml_opt_loss( ggml_opt_context_t opt_ctx); // scalar tensor that contains the loss
152 GGML_API struct ggml_tensor * ggml_opt_pred( ggml_opt_context_t opt_ctx); // predictions made by outputs
153 GGML_API struct ggml_tensor * ggml_opt_ncorrect(ggml_opt_context_t opt_ctx); // number of matching predictions between outputs and labels
154
155 // get the gradient accumulator for a node from the forward graph
156 GGML_API struct ggml_tensor * ggml_opt_grad_acc(ggml_opt_context_t opt_ctx, struct ggml_tensor * node);
157
158 GGML_API enum ggml_opt_optimizer_type ggml_opt_context_optimizer_type(ggml_opt_context_t); //TODO consistent naming scheme
159
160 GGML_API const char * ggml_opt_optimizer_name(enum ggml_opt_optimizer_type);
161
162 // ====== Optimization Result ======
163
164 GGML_API ggml_opt_result_t ggml_opt_result_init(void);
165 GGML_API void ggml_opt_result_free(ggml_opt_result_t result);
166 GGML_API void ggml_opt_result_reset(ggml_opt_result_t result);
167
168 // get data from result, uncertainties are optional and can be ignored by passing NULL
169 GGML_API void ggml_opt_result_ndata( ggml_opt_result_t result, int64_t * ndata); // writes 1 value, number of datapoints
170 GGML_API void ggml_opt_result_loss( ggml_opt_result_t result, double * loss, double * unc); // writes 1 value
171 GGML_API void ggml_opt_result_pred( ggml_opt_result_t result, int32_t * pred); // writes ndata values
172 GGML_API void ggml_opt_result_accuracy(ggml_opt_result_t result, double * accuracy, double * unc); // writes 1 value
173
174 // ====== Computation ======
175
176 // if not using static graphs, this function must be called prior to ggml_opt_alloc
177 GGML_API void ggml_opt_prepare_alloc(
178 ggml_opt_context_t opt_ctx,
179 struct ggml_context * ctx_compute,
180 struct ggml_cgraph * gf,
181 struct ggml_tensor * inputs,
182 struct ggml_tensor * outputs);
183
184 // allocate the next graph for evaluation, either forward or forward + backward
185 // must be called exactly once prior to calling ggml_opt_eval
186 GGML_API void ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward);
187
188 // do forward pass, increment result if not NULL, do backward pass if allocated
189 GGML_API void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result);
190
191 // ############################################################################
192 // ## The high-level functions start here. They do not depend on any private ##
193 // ## functions or structs and can be copied to and adapted for user code. ##
194 // ############################################################################
195
196 // ====== Intended Usage ======
197 //
198 // 1. Select the appropriate loss for your problem.
199 // 2. Create a dataset and set the data for the "data" tensor. Also set the "labels" tensor if your loss needs them.
200 // Setting the shard size to 1 will be fine, it's the granularity with which data is shuffled/loaded (bigger values are faster).
201 // 3. Create a GGML graph for your model with no_alloc == true. Use two separate contexts for the tensors.
202 // The first context should contain the model parameters and inputs and be allocated statically in user code.
203 // The second context should contain all other tensors and will be (re)allocated automatically.
204 // Due to this automated allocation the data of the second context is not defined when accessed in user code.
205 // Note that the second dimension of the inputs/outputs are interpreted as the number of datapoints in those tensors.
206 // 4. Call ggml_opt_fit. If you need more control you can use ggml_opt_epoch instead.
207
208 // signature for a callback while evaluating opt_ctx on dataset, called after an evaluation
209 typedef void (*ggml_opt_epoch_callback)(
210 bool train, // true after training evaluation, false after validation evaluation
211 ggml_opt_context_t opt_ctx,
212 ggml_opt_dataset_t dataset,
213 ggml_opt_result_t result, // result associated with the dataset subsection
214 int64_t ibatch, // number of batches that have been evaluated so far
215 int64_t ibatch_max, // total number of batches in this dataset subsection
216 int64_t t_start_us); // time at which the evaluation on the dataset subsection was started
217
218 // do training on front of dataset, do evaluation only on back of dataset
219 GGML_API void ggml_opt_epoch(
220 ggml_opt_context_t opt_ctx,
221 ggml_opt_dataset_t dataset,
222 ggml_opt_result_t result_train, // result to increment during training, ignored if NULL
223 ggml_opt_result_t result_eval, // result to increment during evaluation, ignored if NULL
224 int64_t idata_split, // data index at which to split training and evaluation
225 ggml_opt_epoch_callback callback_train,
226 ggml_opt_epoch_callback callback_eval);
227
228 // callback that prints a progress bar on stderr
229 GGML_API void ggml_opt_epoch_callback_progress_bar(
230 bool train,
231 ggml_opt_context_t opt_ctx,
232 ggml_opt_dataset_t dataset,
233 ggml_opt_result_t result,
234 int64_t ibatch,
235 int64_t ibatch_max,
236 int64_t t_start_us);
237
238 // fit model defined by inputs and outputs to dataset
239 GGML_API void ggml_opt_fit(
240 ggml_backend_sched_t backend_sched, // backend scheduler for constructing the compute graphs
241 struct ggml_context * ctx_compute, // context with temporarily allocated tensors to calculate the outputs
242 struct ggml_tensor * inputs, // input tensor with shape [ne_datapoint, ndata_batch]
243 struct ggml_tensor * outputs, // output tensor, must have shape [ne_label, ndata_batch] if labels are used
244 ggml_opt_dataset_t dataset, // dataset with data and optionally also labels
245 enum ggml_opt_loss_type loss_type, // loss to minimize
246 enum ggml_opt_optimizer_type optimizer, // sgd or adamw
247 ggml_opt_get_optimizer_params get_opt_pars, // callback to get optimizer params, userdata is pointer to epoch (of type int64_t)
248 int64_t nepoch, // how many times the dataset should be iterated over
249 int64_t nbatch_logical, // datapoints optimizer step, must be a multiple of ndata_batch in inputs/outputs
250 float val_split, // fraction of the dataset to use for validation, must be in [0.0f, 1.0f)
251 bool silent); // whether or not info prints to stderr should be suppressed
252
253
254#ifdef __cplusplus
255}
256#endif