15 lines
406 B
Python
15 lines
406 B
Python
from mflux.config.config import Config
|
|
|
|
|
|
class WeightUtil:
|
|
@staticmethod
|
|
def flatten(params):
|
|
return [(k, v) for p in params for (k, v) in p]
|
|
|
|
@staticmethod
|
|
def reshape_weights(key, value):
|
|
if len(value.shape) == 4:
|
|
value = value.transpose(0, 2, 3, 1)
|
|
value = value.reshape(-1).reshape(value.shape).astype(Config.precision)
|
|
return [(key, value)]
|