blob: 7216c36f7ff4b5c36bc68799c7421740954ec730 (
plain)
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
|
#pragma once
#include "common/config.h"
#include "data/batch_generator.h"
#include "data/corpus_sqlite.h"
#include "models/model_task.h"
#include "training/scheduler.h"
#include "training/validator.h"
namespace marian {
template <class ModelWrapper>
class Train : public ModelTask {
private:
Ptr<Config> options_;
public:
Train(Ptr<Config> options) : options_(options) {}
void run() {
using namespace data;
Ptr<CorpusBase> dataset;
if(!options_->get<std::string>("sqlite").empty())
dataset = New<CorpusSQLite>(options_);
else
dataset = New<Corpus>(options_);
dataset->prepare();
Ptr<BatchStats> stats;
if(options_->get<bool>("mini-batch-fit")) {
LOG(info,
"[batching] Collecting statistics for batch fitting with step size "
"{}",
options_->get<size_t>("mini-batch-fit-step"));
// @TODO, better fake batch with vocabulary
auto model = New<ModelWrapper>(options_);
THREAD_GUARD(stats = model->collectStats());
LOG(info, "[batching] Done");
}
auto trainState = New<TrainingState>(options_->get<float>("learn-rate"));
auto scheduler = New<Scheduler>(options_, trainState);
if((options_->has("valid-sets") || options_->has("valid-script-path"))
&& options_->get<size_t>("valid-freq") > 0) {
for(auto validator : Validators(dataset->getVocabs(), options_))
scheduler->addValidator(validator);
}
auto batchGenerator = New<CorpusBatchGenerator>(dataset, options_, stats);
scheduler->registerTrainingObserver(batchGenerator);
auto model = New<ModelWrapper>(options_);
model->setScheduler(scheduler);
model->load();
// @TODO: shuffle_ as a private attribute in BG
auto shuffle = !options_->get<bool>("no-shuffle");
bool restored = !options_->get<bool>("no-restore-corpus")
&& batchGenerator->restore(trainState, shuffle);
scheduler->started();
while(scheduler->keepGoing()) {
if(!restored)
batchGenerator->prepare(shuffle);
restored = false;
while(*batchGenerator && scheduler->keepGoing()) {
auto batch = batchGenerator->next();
model->update(batch);
}
if(scheduler->keepGoing())
scheduler->increaseEpoch();
}
scheduler->finished();
model->wait();
model->save(true);
}
};
}
|
