def train_and_save_mnist_model(experiment_path, clean_train, triggered_train, clean_test, triggered_test,
model_save_dir, parallel, use_gpu):
logger.info("Training Model...")
def img_transform(x):
return x.unsqueeze(0)
logging_params = {
'num_batches_per_logmsg': 500,
'tensorboard_output_dir': 'tensorboard_dir/',
'experiment_name': 'badnets',
'num_batches_per_metrics': 500,
'num_batches_ver_val_dataset_metrics': None,
'num_epochs_per_metric': 10
}
logging_cfg = tpmc.ReportingConfig(num_batches_per_logmsg=logging_params['num_batches_per_logmsg'],
tensorboard_output_dir=logging_params['tensorboard_output_dir'],
experiment_name=logging_params['experiment_name'],
num_batches_per_metrics=logging_params['num_batches_per_metrics'],
num_batches_ver_val_dataset_metrics=logging_params[
'num_batches_ver_val_dataset_metrics'],
num_epochs_per_metric=logging_params['num_epochs_per_metric'])
# Train clean model to use as a base for triggered model
device = torch.device('cuda' if use_gpu else 'cpu')
data_obj = tpm_tdm.DataManager(experiment_path,
[clean_train, triggered_train],
clean_test,
triggered_test_file=triggered_test,
data_transform=img_transform,
shuffle_train=True)
class MyArchFactory(tpm_af.ArchitectureFactory):
def new_architecture(self):
return tpma.BadNetExample()
clean_training_cfg = tpmc.TrainingConfig(device=device,
epochs=10,
batch_size=100,
lr=1e-4)
clean_optim_cfg = tpmc.DefaultOptimizerConfig(clean_training_cfg, logging_cfg)
clean_optim = tpm_do.DefaultOptimizer(clean_optim_cfg)
triggered_training_cfg = tpmc.TrainingConfig(device=device,
epochs=200,
batch_size=15,
lr=1e-4)
triggered_optim_cfg = tpmc.DefaultOptimizerConfig(triggered_training_cfg, logging_cfg)
triggered_optim = tpm_do.DefaultOptimizer(triggered_optim_cfg)
optims = [clean_optim, triggered_optim]
model_filename = 'BadNets_0.2_poison_sequential.pt'
cfg = tpmc.RunnerConfig(MyArchFactory(), data_obj, optimizer=optims, model_save_dir=model_save_dir,
stats_save_dir=model_save_dir,
filename=model_filename,
parallel=parallel)
runner = tpmr.Runner(cfg, {'script': 'gen_and_train_mnist_sequential.py'})
runner.run()
def train_and_save_mnist_model(experiment_path, triggered_train, clean_test,
triggered_test, model_save_dir, parallel,
use_gpu):
logger.info("Training Model...")
def img_transform(x):
return x.unsqueeze(0)
logging_params = {
'num_batches_per_logmsg': 500,
'tensorboard_output_dir': 'tensorboard_dir/',
'experiment_name': 'badnets',
'num_batches_per_metrics': 500,
'num_epochs_per_metric': 10
}
logging_cfg = tpmc.ReportingConfig(
num_batches_per_logmsg=logging_params['num_batches_per_logmsg'],
tensorboard_output_dir=logging_params['tensorboard_output_dir'],
experiment_name=logging_params['experiment_name'],
num_batches_per_metrics=logging_params['num_batches_per_metrics'],
num_epochs_per_metric=logging_params['num_epochs_per_metric'])
# Train clean model to use as a base for triggered model
device = torch.device('cuda' if use_gpu else 'cpu')
num_avail_cpus = multiprocessing.cpu_count()
num_cpus_to_use = int(.8 * num_avail_cpus)
data_obj = tpm_tdm.DataManager(
experiment_path,
triggered_train,
clean_test,
triggered_test_file=triggered_test,
train_data_transform=img_transform,
test_data_transform=img_transform,
shuffle_train=True,
train_dataloader_kwargs={'num_workers': num_cpus_to_use})
class MyArchFactory(tpm_af.ArchitectureFactory):
def new_architecture(self):
return tpma.ModdedLeNet5Net()
training_cfg = tpmc.TrainingConfig(
device=device,
epochs=300,
batch_size=20,
lr=1e-4,
early_stopping=tpmc.EarlyStoppingConfig())
optim_cfg = tpmc.DefaultOptimizerConfig(training_cfg, logging_cfg)
optim = tpm_do.DefaultOptimizer(optim_cfg)
model_filename = 'ModdedLeNet5_0.2_poison.pt'
cfg = tpmc.RunnerConfig(MyArchFactory(),
data_obj,
optimizer=optim,
model_save_dir=model_save_dir,
stats_save_dir=model_save_dir,
filename=model_filename,
parallel=parallel)
runner = tpmr.Runner(cfg, {'script': 'gen_and_train_mnist.py'})
runner.run()
training_params = tpmc.TrainingConfig(
device=device,
epochs=a.num_epochs,
batch_size=32,
lr=0.001,
optim='adam',
objective='cross_entropy_loss',
early_stopping=early_stopping_argin,
train_val_split=a.train_val_split)
reporting_params = tpmc.ReportingConfig(
num_batches_per_logmsg=500,
num_epochs_per_metric=1,
num_batches_per_metrics=default_nbpvdm,
tensorboard_output_dir=a.tensorboard_dir,
experiment_name=experiment_cfg['name'])
optimizer_cfg = tpmc.DefaultOptimizerConfig(training_params,
reporting_params)
cfg = tpmc.ModelGeneratorConfig(arch,
data_obj,
model_save_dir,
stats_save_dir,
num_models,
optimizer=optimizer_cfg,
experiment_cfg=experiment_cfg,
parallel=True)
# may also provide lists of run_ids or filenames are arguments to ModelGeneratorConfig to have more control
# of saved model file names; see RunnerConfig and ModelGeneratorConfig for more information
modelgen_cfgs.append(cfg)
model_generator = mg.ModelGenerator(modelgen_cfgs)
def test_default_optimizer_config_copy(self):
o1 = tpmc.DefaultOptimizerConfig()
o2 = tpmc.DefaultOptimizerConfig()
self.assertEqual(o1, o2)