def make_model(args):
if args.model == 'fcn_mobilenetv2':
# Convolutional neural netorks
global_model = fcn_mobilenetv2(num_classes=args.num_classes,
aux_loss=bool(args.aux_lr))
if args.no_dropout:
global_model.classifier[3].p = 0
global_model.aux_classifier[3].p = 0
elif args.model == 'deeplabv3_mobilenetv2':
global_model = deeplabv3_mobilenetv2(num_classes=args.num_classes,
aux_loss=bool(args.aux_lr))
if args.no_dropout:
global_model.classifier[3].p = 0
global_model.aux_classifier[3].p = 0
elif args.model == 'deeplabv3_mobilenetv3':
global_model = deeplabv3_mobilenet_v3_large(
num_classes=args.num_classes,
aux_loss=bool(args.aux_lr),
pretrained=args.pretrained)
# fix batchnorm channels divisible by 8
in_channels = global_model.aux_classifier[0].in_channels
global_model.aux_classifier = FCNHead(in_channels, args.num_classes)
if args.no_dropout:
global_model.classifier[0].project[3].p = 0
elif args.model == 'lraspp_mobilenetv3':
global_model = lraspp_mobilenet_v3_large(num_classes=args.num_classes,
pretrained=args.pretrained)
# no aux classifier, no dropout layer
global_model.aux_classifier = None
#resnet for test only as too many params
elif args.model == 'fcn_resnet50':
global_model = fcn_resnet50(num_classes=args.num_classes,
pretrained=True)
else:
exit('Error: unrecognized model')
if args.activation == 'tanh': # test tanh for DP-SGD
global_model = convert_relu_tanh(global_model)
if args.freeze_backbone: # test for DP-SGD
for p in global_model.backbone.parameters():
p.requires_grad = False
# change model architecutre from batch_norm to group_norm for DP
if args.dp:
global_model = convert_batchnorm_modules(global_model)
inspector = DPModelInspector()
assert inspector.validate(global_model) == True
return global_model
def test_dp_model_inspector_example(self):
# IMPORTANT: When changing this code you also need to update
# the docstring for opacus.dp_model_inspector.DPModelInspector.validate()
inspector = DPModelInspector()
valid_model = nn.Linear(16, 32)
is_valid = inspector.validate(valid_model)
self.assertTrue(is_valid)
invalid_model = nn.BatchNorm1d(2)
with self.assertRaises(IncompatibleModuleException):
is_valid = inspector.validate(invalid_model)
def get_trained_model(dataset, dp):
classes, trainloader, testloader, trainset, testset = get_test_train_loaders(
dataset)
device = torch.device("cuda")
net = torchvision.models.alexnet(num_classes=len(classes)).to(device)
if dp:
net = module_modification.convert_batchnorm_modules(net)
inspector = DPModelInspector()
print(f"Model valid: {inspector.validate(net)}")
print(f"Model trained DP: {dp}")
net = net.to(device)
if dp:
PATH = './trained_models/' + dataset + '_dp' + '.pth'
else:
PATH = './trained_models/' + dataset + '.pth'
# PATH='./trained_models/mnist_0_931.pth'
# PATH='./trained_models/mnist_0_dp_316.pth'
print(len(classes))
net.load_state_dict(torch.load(PATH + "",
map_location=torch.device('cpu')))
return net, classes, trainloader, testloader, trainset, testset
def check_dp(model):
inspector = DPModelInspector()
inspector.validate(model)
def setUp(self):
self.validator = DPModelInspector()
elif args.activation == 'tanh':
global_model = CNNCifar10Tanh()
global_model.to(device)
summary(global_model, input_size=(3, 32, 32), device=device)
elif args.dataset == 'dr':
global_model = models.resnet50(num_classes=100)
global_model.to(device)
summary(global_model, input_size=(3, 32, 32), device=device)
else:
exit('Error: unrecognized model')
############# Common ###################
######### DP Model Compatibility #######
if args.withDP:
try:
inspector = DPModelInspector()
inspector.validate(global_model)
print("Model's already Valid!\n")
except:
global_model = module_modification.convert_batchnorm_modules(
global_model)
inspector = DPModelInspector()
print(f"Is the model valid? {inspector.validate(global_model)}")
print("Model is convereted to be Valid!\n")
######### DP Model Compatibility #######
######### Local Models and Optimizers #############
local_models = []
local_optimizers = []
local_privacy_engine = []
class LayersGradTest(unittest.TestCase):
def setUp(self):
self.validator = DPModelInspector()
def _reset_seeds(self):
torch.manual_seed(1337)
torch.cuda.manual_seed(1337)
def _run_once(self, layer, criterion, *args):
self._reset_seeds()
layer.zero_grad()
output = layer(*args)
if isinstance(output, tuple):
output = output[0]
output = output.squeeze()
y = torch.zeros_like(output)
loss = criterion(output, y)
loss.backward()
def _check_one_layer(self, layer, *args, **kwargs):
self._check_one_layer_with_criterion(
layer, nn.L1Loss(reduction="mean"), *args, **kwargs
)
self._check_one_layer_with_criterion(
layer, nn.L1Loss(reduction="sum"), *args, **kwargs
)
def _check_one_layer_with_criterion(self, layer, criterion, *args, **kwargs):
self.validator.validate(layer)
for name, param in layer.named_parameters():
if ("weight" in name) or ("bias" in name):
nn.init.uniform_(param, -1.0, 1.0)
# run without DP
self._run_once(layer, criterion, *args)
vanilla_run_grads = [
(name, p.grad.detach())
for (name, p) in layer.named_parameters()
if p.requires_grad
]
# run with DP
clipper = PerSampleGradientClipper(
layer,
ConstantFlatClipper(1e9),
batch_first=kwargs.get("batch_first", True),
loss_reduction=criterion.reduction,
)
self._run_once(layer, criterion, *args)
for param_name, param in layer.named_parameters():
if param.requires_grad:
self.assertTrue(
hasattr(param, "grad_sample"),
f"Per-sample gradients haven't been computed for {param_name}",
)
clipper.clip_and_accumulate()
clipper.pre_step()
private_run_grads = [
(name, p.grad.detach())
for (name, p) in layer.named_parameters()
if p.requires_grad
]
# compare
for (vanilla_name, vanilla_grad), (private_name, private_grad) in zip(
vanilla_run_grads, private_run_grads
):
assert vanilla_name == private_name
self.assertTrue(
torch.allclose(vanilla_grad, private_grad, atol=10e-5, rtol=10e-3),
f"Gradient mismatch. Parameter: {layer}.{vanilla_name}, loss: {criterion.reduction}",
)
clipper.close()
def test_conv1d(self):
x = torch.randn(64, 16, 24)
layer = nn.Conv1d(16, 32, 3, 1)
self._check_one_layer(layer, x)
def test_conv2d(self):
x = torch.randn(64, 16, 24, 24)
layer = nn.Conv2d(16, 32, 3, 1)
self._check_one_layer(layer, x)
def test_linear(self):
self._check_one_layer(nn.Linear(8, 4), torch.randn(16, 8))
self._check_one_layer(nn.Linear(8, 4), torch.randn(16, 8, 8))
def test_layernorm(self):
x = torch.randn(64, 16, 24, 24)
self._check_one_layer(nn.LayerNorm(24), x)
self._check_one_layer(nn.LayerNorm((24, 24)), x)
self._check_one_layer(nn.LayerNorm((16, 24, 24)), x)
def test_groupnorm(self):
self._check_one_layer(nn.GroupNorm(4, 16), torch.randn(64, 16, 10))
self._check_one_layer(nn.GroupNorm(4, 16), torch.randn(64, 16, 10, 9))
self._check_one_layer(nn.GroupNorm(4, 16), torch.randn(64, 16, 10, 9, 8))
def test_instancenorm(self):
self._check_one_layer(
nn.InstanceNorm1d(16, affine=True), torch.randn(64, 16, 10)
)
self._check_one_layer(
nn.InstanceNorm2d(16, affine=True), torch.randn(64, 16, 10, 9)
)
self._check_one_layer(
nn.InstanceNorm3d(16, affine=True), torch.randn(64, 16, 10, 9, 8)
)
def test_sequence_bias(self):
x = torch.randn(4, 3, 2)
layer = SequenceBias(2)
self._check_one_layer(layer, x, batch_first=False)
def test_multihead_attention(self):
x = torch.randn(16, 24, 32)
layer = DPMultiheadAttention(32, 1)
self._check_one_layer(layer, x, x, x, batch_first=False)
layer = DPMultiheadAttention(32, 1, bias=True, add_bias_kv=True, dropout=0.05)
self._check_one_layer(layer, x, x, x, batch_first=False)
layer = DPMultiheadAttention(32, 1, bias=True, add_bias_kv=True)
self._check_one_layer(layer, x, x, x, batch_first=False)
layer = DPMultiheadAttention(
32, 1, bias=True, add_bias_kv=True, add_zero_attn=True
)
self._check_one_layer(layer, x, x, x, batch_first=False)
q = torch.randn(16, 24, 32)
k = torch.randn(20, 24, 28)
v = torch.randn(20, 24, 28)
layer = DPMultiheadAttention(
32, 1, bias=True, add_bias_kv=True, add_zero_attn=True, kdim=28, vdim=28
)
self._check_one_layer(layer, q, k, v, batch_first=False)
def test_embedding(self):
layer = nn.Embedding(256, 100)
x0 = torch.randint(0, 255, (16, 8, 4)).long()
x1 = torch.randint(0, 255, (16, 8)).long()
x2 = torch.randint(0, 255, (64,)).long()
self._check_one_layer(layer, x0)
self._check_one_layer(layer, x1)
self._check_one_layer(layer, x2)
def test_lstm_batch_first(self):
# input size : 25 output size : 12 minibatch : 30 sequence length : 20
# Test batch_first=True case
layer = DPLSTM(25, 12, 1, batch_first=True)
x = torch.randn(30, 20, 25)
self._check_one_layer(layer, x, batch_first=True)
def test_lstm_batch_second(self):
# input size : 25 output size : 12 minibatch : 30 sequence length : 20
# Test batch_first=False case
layer = DPLSTM(25, 12, 1, batch_first=False)
x = torch.randn(20, 30, 25)
self._check_one_layer(layer, x, batch_first=False)
validation_batches = ceil(len(dataloader) / args.validations_per_epoch)
print(f"Using sigma={args.sigma} and C={args.max_grad_norm}")
if args.delta > 0:
privacy_engine = PrivacyEngine(
model,
batch_size=batch_size,
sample_size=len(train_dataset),
alphas=[1 + x / 10.0 for x in range(1, 100)] + list(range(12, 64)),
noise_multiplier=args.sigma,
max_grad_norm=args.max_grad_norm,
)
model = module_modification.convert_batchnorm_modules(model).cuda()
inspector = DPModelInspector()
print(inspector.validate(model), "--------------")
privacy_engine.attach(optimizer)
start_epoch = 0
# a list of validation IWAE estimates
validation_iwae = []
# a list of running variational lower bounds on the train set
train_vlb = []
# the length of two lists above is the same because the new
# values are inserted into them at the validation checkpoints only
# load the last checkpoint, if it exists
if exists(join(args.model_dir, 'last_checkpoint_{}.tar'.format(args.exp))):
# if exists(join(args.model_dir, args.ckpt)):
