PointNetTransfer/base_classifier.py at master · mati3230/PointNetTransfer · GitHub

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
from abc import ABC, abstractmethod
from utils import mkdir, file_exists
import numpy as np


class BaseClassifier(ABC):
    def __init__(
            self,
            name,
            n_classes,
            seed=None,
            trainable=True,
            check_numerics=False,
            initializer="glorot_uniform",
            trainable_net=True):
        super().__init__()
        self.name = name
        self.seed = seed
        self.check_numerics = check_numerics
        self.n_classes = n_classes
        self.trainable = trainable
        self.trainable_net = trainable_net
        # self.ignore = ["moving_mean", "moving_variance"]
        self.init_net(
            name=name,
            seed=seed,
            trainable=trainable_net,
            check_numerics=check_numerics,
            initializer=initializer)
        self.init_variables(
            name=name,
            n_classes=n_classes,
            trainable=trainable,
            seed=seed,
            initializer=initializer)

    @abstractmethod
    def __call__(self, obs, training):
        pass

    @abstractmethod
    def init_variables(
            self,
            name,
            n_classes,
            trainable=True,
            seed=None,
            initializer="glorot_uniform"):
        pass

    @abstractmethod
    def init_net(
            self,
            name,
            seed=None,
            trainable=True,
            check_numerics=False,
            initializer="glorot_uniform"):
        pass

    @abstractmethod
    def get_vars(self, net_only=False, head_only=False, with_non_trainable=False):
        pass

    @abstractmethod
    def reset(self):
        pass

    def save(self, directory, filename, net_only=False):
        mkdir(directory)
        vars_ = self.get_vars(net_only=net_only)
        if len(vars_) == 0:
            raise Exception("At least one variable is expected")
        var_dict = {}
        for var_ in vars_:
            #print(str(var_.name))
            var_dict[str(var_.name)] = np.array(var_.value())
        np.savez(directory + "/" + filename + ".npz", **var_dict)

    def load(self, directory, filename, net_only=False):
        filepath = directory + "/" + filename + ".npz"
        if not file_exists(filepath):
            raise Exception("File path '" + filepath + "' does not exist")
        model_data = np.load(filepath, allow_pickle=True)
        vars_ = self.get_vars(net_only=net_only)
        if net_only:
            keys = list(model_data.keys())
            for i in range(len(vars_)):
                var_name = vars_[i].name
                """
                tmp = var_name.split("/")[-1]
                tmp = tmp.split(":")[0]
                if tmp in self.ignore:
                    continue
                """
                if var_name not in keys:
                    print(keys)
                    raise Exception("Got no variable with the name " + var_name)
                model_var = model_data[var_name]
                vars_[i].assign(model_var)
        else:
            if len(vars_) != len(model_data):
                keys = list(model_data.keys())
                print("Expected:", len(vars_), "layer; Got:", len(model_data), "layer, file:", filepath)
                if len(vars_) == 0 or len(model_data) == 0:
                    raise Exception("You have to apply a prediction with, e.g., random data to initialize the weights of the network.")
                for i in range(min(len(vars_), len(model_data))):
                    print(vars_[i].name, "\t", keys[i])
                print("Expected:")
                for i in range(len(vars_)):
                    print(vars_[i].name)
                raise Exception("data mismatch")
            i = 0
            for key, value in model_data.items():
                varname = str(vars_[i].name)
                if np.isnan(value).any():
                    raise Exception("loaded value is NaN")
                if key != varname:
                    raise Exception(
                        "Variable names mismatch: " + key + ", " + varname)
                vars_[i].assign(value)
                i += 1