mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
b2cf0fad15
Summary: First part of adding half-floats support to DPER 2.0. Let's add an option use_half_floats to enable converting some weights of the model from fp32 to fp16 before saving it to predictor models parts. For now it's for SparseLookup layer's embeddings. All conversion is done after training is finished and saved models are ready to be used on remote predictors as-is (they will be stored compacted in memory). New fp16 blobs are saved to the model instead of original ones, under the same names, so we don't modify MetaNetDef at all. Next steps: 1) support on delivery side -- operators working with these blobs should support both float and float16 input types 2) benchmark performance to make sure there is no regression a) of serialization b) of delivery 3) support realtime training (I'm thinking about adding new pre-publishing net which will be executed each time the realtime trainer stops to publish a new snapshot) Depends on D4567304 Reviewed By: kennyhorror Differential Revision: D4571710 fbshipit-source-id: 19967a17d3bd84878d66e8c0ed8c5342bf38d979
116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
from caffe2.python import schema
|
|
from caffe2.python.layers.tags import TagContext
|
|
|
|
from collections import namedtuple
|
|
import numpy as np
|
|
|
|
# Some types to simplify descriptions of things traveling between ops
|
|
IdList = schema.List(np.int64)
|
|
IdScoreList = schema.Map(np.int64, np.float32)
|
|
|
|
|
|
class InstantiationContext(object):
|
|
"""
|
|
List of contexts where layer could be instantitated
|
|
"""
|
|
TRAINING = 'training'
|
|
PREDICTION = 'prediction'
|
|
|
|
|
|
_LAYER_REGISTRY = {}
|
|
|
|
|
|
def register_layer(name, layer):
|
|
assert name not in _LAYER_REGISTRY, "{0} already exists".format(name)
|
|
_LAYER_REGISTRY[name] = layer
|
|
|
|
|
|
def layer_exists(name):
|
|
return name in _LAYER_REGISTRY
|
|
|
|
|
|
def create_layer(name, *args, **kwargs):
|
|
return _LAYER_REGISTRY[name](*args, **kwargs)
|
|
|
|
|
|
# TODO(amalevich): Modify this to some better struct, something closer to
|
|
# ParameterInfo.
|
|
LayerParameter = namedtuple(
|
|
'LayerParameter', ['parameter', 'optimizer', 'initializer'])
|
|
|
|
|
|
def _is_request_only_scalar(scalar):
|
|
if len(scalar.field_metadata()) == 0:
|
|
return False
|
|
for metadata in scalar.field_metadata():
|
|
if not (metadata and metadata.feature_specs and
|
|
metadata.feature_specs.feature_is_request_only):
|
|
return False
|
|
return True
|
|
|
|
|
|
class ModelLayer(object):
|
|
|
|
def __init__(self, model, prefix, input_record, tags=set(), **kwargs):
|
|
self.name = model.next_layer_name(prefix)
|
|
self.model = model
|
|
self.kwargs = kwargs
|
|
self.input_record = input_record
|
|
self.request_only = True
|
|
if len(input_record.all_scalars()) == 0:
|
|
self.request_only = False
|
|
for scalar in input_record.all_scalars():
|
|
if not _is_request_only_scalar(scalar):
|
|
self.request_only = False
|
|
break
|
|
self.output_schema = None
|
|
self.tags = set(tags)
|
|
self.tags.update(TagContext.current().tags)
|
|
self.params = []
|
|
|
|
def get_type(self):
|
|
return self.__class__.__name__
|
|
|
|
def get_output_schema(self):
|
|
assert self.output_schema is not None, "Schema is not initialized"
|
|
return self.output_schema
|
|
|
|
def get_parameters(self):
|
|
return self.params
|
|
|
|
def get_fp16_compatible_parameters(self):
|
|
"""Return a subset of parameters which can be converted to fp16"""
|
|
return []
|
|
|
|
def get_memory_usage(self):
|
|
return 0
|
|
|
|
def add_operators(self, net, init_net=None,
|
|
context=InstantiationContext.TRAINING):
|
|
if context != InstantiationContext.PREDICTION:
|
|
assert init_net,\
|
|
"Only prediction context can be used without init_net"
|
|
if init_net:
|
|
for param in self.params:
|
|
# TODO(amalevich): Either return back to lambdas, that add all
|
|
# params (looks a bit safer and breaking less abstractions) or
|
|
# extend Net interface to this type of operations better
|
|
init_net._net.op.extend([param.initializer])
|
|
if context == InstantiationContext.TRAINING:
|
|
self.add_train_ops(net)
|
|
else:
|
|
self.add_ops(net)
|
|
|
|
def add_ops(self, net):
|
|
raise NotImplementedError
|
|
|
|
def add_train_ops(self, net):
|
|
# Default train layer implementation is completely matching predict
|
|
# layer implementation.
|
|
self.add_ops(net)
|