| |
| |
|
|
|
|
|
|
|
|
|
|
| import abc |
|
|
|
|
| class SamplingTrainableMixin(metaclass=abc.ABCMeta): |
|
|
| def __init__(self, *args, **kwargs): |
| super(SamplingTrainableMixin, self).__init__(*args, **kwargs) |
| self._train_param_blobs = None |
| self._train_param_blobs_frozen = False |
|
|
| @property |
| @abc.abstractmethod |
| def param_blobs(self): |
| """ |
| List of parameter blobs for prediction net |
| """ |
| pass |
|
|
| @property |
| def train_param_blobs(self): |
| """ |
| If train_param_blobs is not set before used, default to param_blobs |
| """ |
| if self._train_param_blobs is None: |
| self.train_param_blobs = self.param_blobs |
| return self._train_param_blobs |
|
|
| @train_param_blobs.setter |
| def train_param_blobs(self, blobs): |
| assert not self._train_param_blobs_frozen |
| assert blobs is not None |
| self._train_param_blobs_frozen = True |
| self._train_param_blobs = blobs |
|
|
| @abc.abstractmethod |
| def _add_ops(self, net, param_blobs): |
| """ |
| Add ops to the given net, using the given param_blobs |
| """ |
| pass |
|
|
| def add_ops(self, net): |
| self._add_ops(net, self.param_blobs) |
|
|
| def add_train_ops(self, net): |
| self._add_ops(net, self.train_param_blobs) |
|
|
