| | |
| | |
| |
|
| |
|
| |
|
| |
|
| |
|
| | from caffe2.python import schema |
| | from caffe2.python.layers.layers import ModelLayer, get_layer_class |
| | from caffe2.python.layers.sampling_trainable_mixin import SamplingTrainableMixin |
| |
|
| |
|
| | class SamplingTrain(ModelLayer): |
| | def __init__( |
| | self, |
| | model, |
| | input_record, |
| | prediction_layer, |
| | output_dims, |
| | subtract_log_odd=True, |
| | name='sampling_train', |
| | **kwargs |
| | ): |
| | super(SamplingTrain, self).__init__( |
| | model, name, input_record, **kwargs |
| | ) |
| |
|
| | layer_class = get_layer_class(prediction_layer) |
| | assert issubclass(layer_class, SamplingTrainableMixin) |
| |
|
| | assert 'indices' in input_record |
| | assert isinstance(input_record.indices, schema.Scalar),\ |
| | "input_record.indices is expected to be a schema.Scalar" |
| | assert 'input' in input_record |
| |
|
| | self.subtract_log_odd = subtract_log_odd |
| | if self.subtract_log_odd: |
| | assert 'sampling_prob' in input_record |
| |
|
| | self._prediction_layer = layer_class( |
| | model, |
| | input_record.input, |
| | output_dims=output_dims, |
| | **kwargs |
| | ) |
| |
|
| | self._prediction_layer.train_param_blobs = [ |
| | model.net.NextBlob(str(blob) + '_sampled') |
| | for blob in self._prediction_layer.param_blobs |
| | ] |
| |
|
| | self.params = self._prediction_layer.params |
| |
|
| | self.output_schema = self._prediction_layer.output_schema |
| |
|
| | def add_ops(self, net): |
| | self._prediction_layer.add_ops(net) |
| |
|
| | def add_train_ops(self, net): |
| | for full_blob, sampled_blob in zip( |
| | self._prediction_layer.param_blobs, |
| | self._prediction_layer.train_param_blobs |
| | ): |
| | net.Gather([full_blob, self.input_record.indices()], sampled_blob) |
| | self._prediction_layer.add_train_ops(net) |
| | if not self.subtract_log_odd: |
| | return |
| | log_q = net.Log(self.input_record.sampling_prob(), |
| | net.NextScopedBlob("log_q")) |
| | net.Sub([self.output_schema(), log_q], self.output_schema(), |
| | broadcast=1, use_grad_hack=1) |
| |
|
