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
|
#pragma once
#include <future>
#include <boost/filesystem.hpp>
#include "training/graph_group.h"
namespace marian {
class SingletonGraph : public GraphGroup {
public:
virtual void setScheduler(Ptr<Scheduler> scheduler);
private:
Ptr<models::ModelBase> builder_;
Ptr<ExpressionGraph> graph_;
Ptr<ExpressionGraph> mvAvgGraph_;
bool mvAvg_{false};
float mvDecay_{1e-4};
void updateMovingAverage(Tensor mvAvgParams, Tensor params, size_t batches);
void execute(Ptr<data::Batch> batch);
public:
SingletonGraph(Ptr<Config> options)
: GraphGroup(options),
mvAvg_{options_->get<float>("exponential-smoothing") > 0},
mvDecay_{options_->get<float>("exponential-smoothing")} {
auto deviceId = options_->getDevices()[0];
graph_ = New<ExpressionGraph>();
graph_->setDevice(deviceId);
graph_->reserveWorkspaceMB(options_->get<size_t>("workspace"));
opt_ = Optimizer(options_);
builder_ = models::from_config(options_);
}
void update(Ptr<data::Batch> batch) { execute(batch); }
void load() {
if(!options_->get<bool>("no-reload")) {
std::string name = options_->get<std::string>("model");
if(boost::filesystem::exists(name)) {
if(scheduler_)
scheduler_->load(name);
builder_->load(graph_, name);
opt_->load(name + ".optimizer.npz", {opt_}, {graph_->getBackend()});
} else if(options_->has("pretrained-model")) {
std::string init = options_->get<std::string>("pretrained-model");
LOG(info,
"Initialize model weights with the pre-trained model {}",
init);
builder_->load(graph_, init, false);
}
}
}
void save(bool final = false) {
auto saveGraph = graph_;
if(mvAvg_)
saveGraph = mvAvgGraph_;
if(final && scheduler_)
scheduler_->validate({saveGraph}, true);
save(saveGraph, final);
}
void save(Ptr<ExpressionGraph> graph, bool final = false) {
std::string name = options_->get<std::string>("model");
if(options_->get<bool>("overwrite")) {
builder_->save(graph_, name, true);
if(scheduler_)
scheduler_->save(name);
} else {
if(!final) {
std::string numberOfBatches
= scheduler_ ? std::to_string(scheduler_->numberOfBatches())
: "unknown";
std::string nameOverwrite = name;
nameOverwrite.replace(
name.size() - 4, 4, ".iter" + numberOfBatches + ".npz");
builder_->save(graph_, nameOverwrite);
}
builder_->save(graph_, name, true);
if(scheduler_)
scheduler_->save(name);
}
size_t totalSize = graph_->params()->vals()->size();
opt_->save(name + ".optimizer.npz", {opt_}, totalSize);
}
Ptr<data::BatchStats> collectStats() {
return builder_->collectStats(graph_);
}
};
}
|
