-
Notifications
You must be signed in to change notification settings - Fork 0
/
tf1dp.py
154 lines (132 loc) · 6.53 KB
/
tf1dp.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
import os
import time
from functools import partial
import tensorflow.compat.v1 as tf
from tensorflow_privacy.privacy.analysis.gdp_accountant import (compute_eps_poisson,
compute_mu_poisson)
from tensorflow_privacy.privacy.analysis.rdp_accountant import (compute_rdp, get_privacy_spent)
from tensorflow_privacy.privacy.optimizers import dp_optimizer_vectorized
import data
import utils
from tf2dp import model_dict
def nn_model_fn(model, loss_fn, args, features, labels, mode):
logits = model()(features['x'])
vector_loss = loss_fn(labels=labels, logits=logits)
scalar_loss = tf.reduce_mean(vector_loss)
if mode == tf.estimator.ModeKeys.TRAIN:
if args.dpsgd:
# Use DP version of GradientDescentOptimizer. Other optimizers are
# available in dp_optimizer. Most optimizers inheriting from
# tf.train.Optimizer should be wrappable in differentially private
# counterparts by calling dp_optimizer.optimizer_from_args().
optimizer = dp_optimizer_vectorized.VectorizedDPSGD(
l2_norm_clip=args.l2_norm_clip,
noise_multiplier=args.noise_multiplier,
num_microbatches=args.microbatches,
learning_rate=args.learning_rate)
opt_loss = vector_loss
else:
optimizer = tf.train.GradientDescentOptimizer(learning_rate=args.learning_rate)
opt_loss = scalar_loss
global_step = tf.train.get_global_step()
train_op = optimizer.minimize(loss=opt_loss, global_step=global_step)
# In the following, we pass the mean of the loss (scalar_loss) rather than
# the vector_loss because tf.estimator requires a scalar loss. This is only
# used for evaluation and debugging by tf.estimator. The actual loss being
# minimized is opt_loss defined above and passed to optimizer.minimize().
return tf.estimator.EstimatorSpec(mode=mode, loss=scalar_loss, train_op=train_op)
elif mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {
'accuracy': tf.metrics.accuracy(labels=labels,
predictions=tf.argmax(input=logits, axis=1))
}
return tf.estimator.EstimatorSpec(mode=mode,
loss=scalar_loss,
eval_metric_ops=eval_metric_ops)
def compute_epsilon(epoch, num_train_eg, args):
"""Computes epsilon value for given hyperparameters."""
steps = epoch * num_train_eg // args.batch_size
if args.noise_multiplier == 0.0:
return float('inf')
orders = [1 + x / 10. for x in range(1, 100)] + list(range(12, 64))
sampling_probability = args.batch_size / num_train_eg
rdp = compute_rdp(q=sampling_probability,
noise_multiplier=args.noise_multiplier,
steps=steps,
orders=orders)
# Delta is set to approximate 1 / (number of training points).
return get_privacy_spent(orders, rdp, target_delta=1e-5)[0]
def main(args):
print(args)
tf.disable_eager_execution()
if args.memory_limit:
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
tf.config.experimental.set_virtual_device_configuration(
physical_devices[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=args.memory_limit)])
assert args.microbatches is None
args.microbatches = args.batch_size
data_fn = data.data_fn_dict[args.experiment][int(args.dummy_data)]
kwargs = {
'max_features': args.max_features,
'max_len': args.max_len,
'format': 'NHWC',
}
if args.dummy_data:
kwargs['num_examples'] = args.batch_size * 2
(train_data, train_labels), _ = data_fn(**kwargs)
num_train_eg = train_data.shape[0]
loss_fn = tf.nn.sparse_softmax_cross_entropy_with_logits
if args.experiment == 'logreg':
loss_fn = lambda labels, logits: tf.nn.sigmoid_cross_entropy_with_logits(
labels=labels, logits=tf.squeeze(logits))
train_labels = train_labels.astype('float32')
model = partial(model_dict[args.experiment],
features=train_data,
max_features=args.max_features,
args=args)
if args.use_xla:
# Not sure which one of these two works, so I'll just use both
assert os.environ['TF_XLA_FLAGS'] == '--tf_xla_auto_jit=2'
session_config = tf.ConfigProto()
session_config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_2
run_config = tf.estimator.RunConfig(session_config=session_config)
print('Using XLA!')
else:
run_config = None
print('NOT using XLA!')
model_obj = tf.estimator.Estimator(model_fn=partial(nn_model_fn, model, loss_fn, args),
config=run_config)
train_input_fn = tf.estimator.inputs.numpy_input_fn(x={'x': train_data},
y=train_labels,
batch_size=args.batch_size,
num_epochs=args.epochs,
shuffle=True)
steps_per_epoch = num_train_eg // args.batch_size
timings = []
for epoch in range(1, args.epochs + 1):
start = time.perf_counter()
model_obj.train(input_fn=train_input_fn, steps=steps_per_epoch)
duration = time.perf_counter() - start
print("Time Taken: ", duration)
timings.append(duration)
if args.dpsgd:
# eps = compute_epsilon(epoch, num_train_eg, args)
# print('For delta=1e-5, the current epsilon is: %.2f' % eps)
print('Trained with DPSGD optimizer')
else:
print('Trained with vanilla non-private SGD optimizer')
if not args.no_save:
utils.save_runtimes(__file__.split('.')[0], args, timings)
else:
print('Not saving!')
print('Done!')
if __name__ == '__main__':
parser = utils.get_parser(model_dict.keys())
parser.add_argument('--memory_limit', default=None, type=int)
parser.add_argument('--xla', dest='use_xla', action='store_true')
parser.add_argument('--no_xla', dest='use_xla', action='store_false')
parser.add_argument('--no_unroll', dest='no_unroll', action='store_true')
args = parser.parse_args()
main(args)