File size: 2,903 Bytes
fcd5579 |
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 |
import pickle
from pathlib import Path
from core import pathex
import numpy as np
from core.leras import nn
tf = nn.tf
class Saveable():
def __init__(self, name=None):
self.name = name
#override
def get_weights(self):
#return tf tensors that should be initialized/loaded/saved
return []
#override
def get_weights_np(self):
weights = self.get_weights()
if len(weights) == 0:
return []
return nn.tf_sess.run (weights)
def set_weights(self, new_weights):
weights = self.get_weights()
if len(weights) != len(new_weights):
raise ValueError ('len of lists mismatch')
tuples = []
for w, new_w in zip(weights, new_weights):
if len(w.shape) != new_w.shape:
new_w = new_w.reshape(w.shape)
tuples.append ( (w, new_w) )
nn.batch_set_value (tuples)
def save_weights(self, filename, force_dtype=None):
d = {}
weights = self.get_weights()
if self.name is None:
raise Exception("name must be defined.")
name = self.name
for w in weights:
w_val = nn.tf_sess.run (w).copy()
w_name_split = w.name.split('/', 1)
if name != w_name_split[0]:
raise Exception("weight first name != Saveable.name")
if force_dtype is not None:
w_val = w_val.astype(force_dtype)
d[ w_name_split[1] ] = w_val
d_dumped = pickle.dumps (d, 4)
pathex.write_bytes_safe ( Path(filename), d_dumped )
def load_weights(self, filename):
"""
returns True if file exists
"""
filepath = Path(filename)
if filepath.exists():
result = True
d_dumped = filepath.read_bytes()
d = pickle.loads(d_dumped)
else:
return False
weights = self.get_weights()
if self.name is None:
raise Exception("name must be defined.")
try:
tuples = []
for w in weights:
w_name_split = w.name.split('/')
if self.name != w_name_split[0]:
raise Exception("weight first name != Saveable.name")
sub_w_name = "/".join(w_name_split[1:])
w_val = d.get(sub_w_name, None)
if w_val is None:
#io.log_err(f"Weight {w.name} was not loaded from file {filename}")
tuples.append ( (w, w.initializer) )
else:
w_val = np.reshape( w_val, w.shape.as_list() )
tuples.append ( (w, w_val) )
nn.batch_set_value(tuples)
except:
return False
return True
def init_weights(self):
nn.init_weights(self.get_weights())
nn.Saveable = Saveable
|