forked from 626_privacy/tensorflow_privacy
4d0ab48c35
The privacy ledger keeps a record of all sampling and query events for analysis post hoc by the privacy accountant. PiperOrigin-RevId: 233094012
91 lines
3 KiB
Python
91 lines
3 KiB
Python
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""A lightweight fixed-sized buffer for maintaining lists.
|
|
"""
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import tensorflow as tf
|
|
|
|
|
|
class TensorBuffer(object):
|
|
"""A lightweight fixed-sized buffer for maintaining lists.
|
|
|
|
The TensorBuffer accumulates tensors of the given shape into a tensor (whose
|
|
rank is one more than that of the given shape) via calls to `append`. The
|
|
current value of the accumulated tensor can be extracted via the property
|
|
`values`.
|
|
"""
|
|
|
|
def __init__(self, max_size, shape, dtype=tf.int32, name=None):
|
|
"""Initializes the TensorBuffer.
|
|
|
|
Args:
|
|
max_size: The maximum size. Attempts to append more than this many rows
|
|
will fail with an exception.
|
|
shape: The shape (as tuple or list) of the tensors to accumulate.
|
|
dtype: The type of the tensors.
|
|
name: A string name for the variable_scope used.
|
|
|
|
Raises:
|
|
ValueError: If the shape is empty (specifies scalar shape).
|
|
"""
|
|
shape = list(shape)
|
|
self._rank = len(shape)
|
|
if not self._rank:
|
|
raise ValueError('Shape cannot be scalar.')
|
|
shape = [max_size] + shape
|
|
|
|
with tf.variable_scope(name):
|
|
self._buffer = tf.Variable(
|
|
initial_value=tf.zeros(shape, dtype),
|
|
trainable=False,
|
|
name='buffer')
|
|
self._size = tf.Variable(
|
|
initial_value=0,
|
|
trainable=False,
|
|
name='size')
|
|
|
|
def append(self, value):
|
|
"""Appends a new tensor to the end of the buffer.
|
|
|
|
Args:
|
|
value: The tensor to append. Must match the shape specified in the
|
|
initializer.
|
|
|
|
Returns:
|
|
An op appending the new tensor to the end of the buffer.
|
|
"""
|
|
with tf.control_dependencies([
|
|
tf.assert_less(
|
|
self._size,
|
|
tf.shape(self._buffer)[0],
|
|
message='Appending past end of TensorBuffer.'),
|
|
tf.assert_equal(
|
|
tf.shape(value),
|
|
tf.shape(self._buffer)[1:],
|
|
message='Appending value of inconsistent shape.')]):
|
|
with tf.control_dependencies(
|
|
[tf.assign(self._buffer[self._size, :], value)]):
|
|
return tf.assign_add(self._size, 1)
|
|
|
|
@property
|
|
def values(self):
|
|
"""Returns the accumulated tensor."""
|
|
begin_value = tf.zeros([self._rank + 1], dtype=tf.int32)
|
|
value_size = tf.concat(
|
|
[[self._size], tf.constant(-1, tf.int32, [self._rank])], 0)
|
|
return tf.slice(self._buffer, begin_value, value_size)
|