ruby/mia_on_model_distillation
1
0
Fork 0
forked from 626_privacy/mia_on_model_distillation
Code Pull requests Activity

Wres: cuda and norm flags

Browse source
This commit is contained in:
Akemi Izuko 2024-11-30 23:54:48 -07:00
parent e4b5998dbb
commit 424cb01a15
Signed by: akemi
GPG key ID: 8DE0764E1809E9FC
2 changed files with 71 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
wresnet-pytorch/src/WideResNet.py
View file

@ -5,7 +5,6 @@ import math
class IndividualBlock1(nn.Module): class IndividualBlock1(nn.Module):
def __init__(self, input_features, output_features, stride, subsample_input=True, increase_filters=True): def __init__(self, input_features, output_features, stride, subsample_input=True, increase_filters=True):
super(IndividualBlock1, self).__init__() super(IndividualBlock1, self).__init__()

36
wresnet-pytorch/src/train.py
View file

@ -10,6 +10,7 @@ from WideResNet import WideResNet
from tqdm import tqdm from tqdm import tqdm
import opacus import opacus
from opacus.validators import ModuleValidator from opacus.validators import ModuleValidator
from opacus.utils.batch_memory_manager import BatchMemoryManager
def set_seed(seed=42): def set_seed(seed=42):
@ -20,15 +21,15 @@ def set_seed(seed=42):
torch.cuda.manual_seed(seed) torch.cuda.manual_seed(seed)
def _train_seed(net, loaders, device, dataset, log=False, checkpoint=False, logfile='', checkpointFile='', epochs=200): def _train_seed(net, loaders, device, dataset, log=False, checkpoint=False, logfile='', checkpointFile='', epochs=200, norm=1.0):
train_loader, test_loader = loaders train_loader, test_loader = loaders
dp_epsilon = 8 dp_epsilon = 8
dp_delta = 1e-5
if dp_epsilon is not None: if dp_epsilon is not None:
print(f"DP epsilon: {dp_epsilon}") print(f"DP epsilon = {dp_epsilon}, delta = {dp_delta}")
#net = ModuleValidator.fix(net, replace_bn_with_in=True) #net = ModuleValidator.fix(net, replace_bn_with_in=True)
net = ModuleValidator.fix(net) net = ModuleValidator.fix(net)
print(net)
ModuleValidator.validate(net, strict=True) ModuleValidator.validate(net, strict=True)
criterion = nn.CrossEntropyLoss() criterion = nn.CrossEntropyLoss()
@ -37,25 +38,33 @@ def _train_seed(net, loaders, device, dataset, log=False, checkpoint=False, logf
best_test_set_accuracy = 0 best_test_set_accuracy = 0
if dp_epsilon is not None:
privacy_engine = opacus.PrivacyEngine() privacy_engine = opacus.PrivacyEngine()
net, optimizer, train_loader = privacy_engine.make_private_with_epsilon( net, optimizer, train_loader = privacy_engine.make_private_with_epsilon(
module=net, module=net,
optimizer=optimizer, optimizer=optimizer,
data_loader=train_loader, data_loader=train_loader,
epochs=epochs, epochs=epochs,
target_epsilon=8, target_epsilon=dp_epsilon,
target_delta=1e-5, target_delta=dp_delta,
max_grad_norm=3.0, max_grad_norm=norm,
) )
print(f"Using sigma={optimizer.noise_multiplier} and C={1.0}") print(f"Using sigma={optimizer.noise_multiplier} and C={1.0}, norm = {norm}")
else:
print("Training without differential privacy")
print(f"Training with {epochs} epochs") print(f"Training with {epochs} epochs")
#for epoch in tqdm(range(epochs)): #for epoch in tqdm(range(epochs)):
with BatchMemoryManager(
data_loader=train_loader,
max_physical_batch_size=1000, # Roughly 12gb vram, uses 9.4
optimizer=optimizer
) as memory_safe_data_loader:
for epoch in range(epochs): for epoch in range(epochs):
net.train() net.train()
#for i, data in tqdm(enumerate(train_loader, 0), leave=False): #for i, data in tqdm(enumerate(train_loader, 0), leave=False):
for i, data in enumerate(train_loader, 0): for i, data in enumerate(memory_safe_data_loader, 0):
inputs, labels = data inputs, labels = data
inputs = inputs.to(device) inputs = inputs.to(device)
labels = labels.to(device) labels = labels.to(device)
@ -70,6 +79,7 @@ def _train_seed(net, loaders, device, dataset, log=False, checkpoint=False, logf
scheduler.step() scheduler.step()
if epoch % 10 == 0 or epoch == epochs - 1:
with torch.no_grad(): with torch.no_grad():
correct = 0 correct = 0
@ -91,7 +101,7 @@ def _train_seed(net, loaders, device, dataset, log=False, checkpoint=False, logf
epoch_accuracy = round(100 * epoch_accuracy, 2) epoch_accuracy = round(100 * epoch_accuracy, 2)
if log: if log:
print('Accuracy at epoch {} is {}%\n'.format(epoch + 1, epoch_accuracy)) print('Accuracy at epoch {} is {}%'.format(epoch + 1, epoch_accuracy))
with open(logfile, 'a') as temp: with open(logfile, 'a') as temp:
temp.write('Accuracy at epoch {} is {}%\n'.format(epoch + 1, epoch_accuracy)) temp.write('Accuracy at epoch {} is {}%\n'.format(epoch + 1, epoch_accuracy))
@ -124,7 +134,9 @@ def train(args):
checkpoint = True if training_configurations.checkpoint.lower() == 'true' else False checkpoint = True if training_configurations.checkpoint.lower() == 'true' else False
loaders = get_loaders(dataset, training_configurations.batch_size) loaders = get_loaders(dataset, training_configurations.batch_size)
if torch.cuda.is_available(): if torch.cuda.is_available() and args.cuda:
device = torch.device(f'cuda:{args.cuda}')
elif torch.cuda.is_available():
device = torch.device('cuda:0') device = torch.device('cuda:0')
else: else:
device = torch.device('cpu') device = torch.device('cpu')
@ -144,7 +156,7 @@ def train(args):
checkpointFile = 'wrn-{}-{}-seed-{}-{}-dict.pth'.format(wrn_depth, wrn_width, dataset, seed) if checkpoint else '' checkpointFile = 'wrn-{}-{}-seed-{}-{}-dict.pth'.format(wrn_depth, wrn_width, dataset, seed) if checkpoint else ''
epochs = training_configurations.epochs epochs = training_configurations.epochs
best_test_set_accuracy = _train_seed(net, loaders, device, dataset, log, checkpoint, logfile, checkpointFile, epochs) best_test_set_accuracy = _train_seed(net, loaders, device, dataset, log, checkpoint, logfile, checkpointFile, epochs, args.norm)
if log: if log:
with open(logfile, 'a') as temp: with open(logfile, 'a') as temp:
@ -168,6 +180,8 @@ if __name__ == '__main__':
parser = argparse.ArgumentParser(description='WideResNet') parser = argparse.ArgumentParser(description='WideResNet')
parser.add_argument('-config', '--config', help='Training Configurations', required=True) parser.add_argument('-config', '--config', help='Training Configurations', required=True)
parser.add_argument('--norm', type=float, help='dpsgd norm clip factor', required=True)
parser.add_argument('--cuda', type=int, help='gpu index', required=False)
args = parser.parse_args() args = parser.parse_args()