CNNCreator.ftl 2.57 KB
Newer Older
1
2
3
import mxnet as mx
import logging
import os
4
5
6
7
8
<#list tc.architecture.streams as stream>
<#if stream.isNetwork()>
from CNNNet_${tc.fullArchitectureName} import Net_${stream?index}
</#if>
</#list>
9
10

class ${tc.fileNameWithoutEnding}:
nilsfreyer's avatar
nilsfreyer committed
11
12
    _model_dir_ = "model/${tc.componentName}/"
    _model_prefix_ = "model"
13

Nicola Gatto's avatar
Nicola Gatto committed
14
15
    def __init__(self):
        self.weight_initializer = mx.init.Normal()
16
        self.networks = {}
17

18
    def load(self, context):
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
        earliestLastEpoch = None

        for i, network in self.networks.items():
            lastEpoch = 0
            param_file = None

            try:
                os.remove(self._model_dir_ + self._model_prefix_ + "_" + str(i) + "_newest-0000.params")
            except OSError:
                pass
            try:
                os.remove(self._model_dir_ + self._model_prefix_ + "_" + str(i) + "_newest-symbol.json")
            except OSError:
                pass

            if os.path.isdir(self._model_dir_):
                for file in os.listdir(self._model_dir_):
                    if ".params" in file and self._model_prefix_ + "_" + str(i) in file:
                        epochStr = file.replace(".params","").replace(self._model_prefix_ + "_" + str(i) + "-","")
                        epoch = int(epochStr)
                        if epoch > lastEpoch:
                            lastEpoch = epoch
                            param_file = file
            if param_file is None:
                earliestLastEpoch = 0
            else:
                logging.info("Loading checkpoint: " + param_file)
                network.load_parameters(self._model_dir_ + param_file)

                if earliestLastEpoch == None or lastEpoch < earliestLastEpoch:
                    earliestLastEpoch = lastEpoch

        return earliestLastEpoch
52
53

    def construct(self, context, data_mean=None, data_std=None):
54
55
56
57
58
59
60
61
<#list tc.architecture.streams as stream>
<#if stream.isNetwork()>
        self.networks[${stream?index}] = Net_${stream?index}(data_mean=data_mean, data_std=data_std)
        self.networks[${stream?index}].collect_params().initialize(self.weight_initializer, ctx=context)
        self.networks[${stream?index}].hybridize()
        self.networks[${stream?index}](<#list stream.getFirstAtomicElements() as input>mx.nd.zeros((1, ${tc.join(input.definition.type.dimensions, ",")},), ctx=context)<#sep>, </#list>)
</#if>
</#list>
Nicola Gatto's avatar
Nicola Gatto committed
62
63
64
65

        if not os.path.exists(self._model_dir_):
            os.makedirs(self._model_dir_)

66
67
        for i, network in self.networks.items():
            network.export(self._model_dir_ + self._model_prefix_ + "_" + str(i), epoch=0)