def main():
parser = argparse.ArgumentParser()
parser.add_argument('--out-dir',
type=str,
required=True,
help='Out directory to save results to')
parser.add_argument('--df-path',
type=str,
required=True,
help='Input dataframe to draw annotations from')
parser.add_argument('--debug', action='store_true')
args = parser.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(exist_ok=True)
s = Store(str(args.out_dir))
# held_outs
s.add_table('out', {'dists': s.OBJECT})
df = load_dfs(path=args.df_path)
dists = model_fit(args.debug, NUM_BETAS, df)
s['out'].append_row({'dists': dists})
print('store located in:', s.path)
print(
'In Beta-Binomial Model Analysis.ipynb: set INPUT_DATA = \'f{s.path}\''
)
def test_L1Ball(algorithm, lr):
# Setup
constraint = chop.constraints.L1Ball(alpha)
prox = constraint.prox
lmo = constraint.lmo
assert (constraint.prox(w) == w).all()
w_t = Variable(torch.zeros_like(w), requires_grad=True)
constraint_oracles = {
stochastic.PGD.name: {
'prox': [prox]
},
stochastic.PGDMadry.name: {
'prox': [prox],
'lmo': [lmo]
},
stochastic.FrankWolfe.name: {
'lmo': [lmo]
},
stochastic.S3CM.name: {
'prox1': [prox],
'prox2': [prox]
}
}
optimizer = algorithm([w_t], **(constraint_oracles[algorithm.name]), lr=lr)
criterion = torch.nn.MSELoss(reduction='mean')
# Logging
store = Store(OUT_DIR)
store.add_table('metadata', {'algorithm': str, 'lr': float})
store['metadata'].append_row({'algorithm': optimizer.name, 'lr': lr})
store.add_table(optimizer.name, {
'func_val': float,
'certificate': float,
'norm(w_t)': float
})
cert = torch.tensor(np.inf)
for ii in range(MAX_ITER):
optimizer.zero_grad()
loss = criterion(X.mv(w_t), y)
loss.backward()
optimizer.step()
try:
cert = next(optimizer.certificate) # only one parameter here
except AttributeError:
cert = torch.tensor(np.nan)
store.log_table_and_tb(
optimizer.name, {
'func_val': loss.item(),
'certificate': cert.item(),
'norm(w_t)': sum(abs(w_t)).item()
})
store[optimizer.name].flush_row()
store.close()
def train_configs(configs: List[Config],
n_seeds: List[int],
targets: List[int],
n_epochs: int,
save_dir_artifacts: str,
save_dir_models: str,
data_aug: bool = False,
save_models: bool = True,
adv_acc: bool = False,
run_time_str: str = None):
loaders_BIN, normalization_function_BIN, label_map_BIN = utils.get_binary_dataset(
batch_size=256, transfer=True, data_aug=data_aug, targets=targets, per_target=128, random=False
)
train_loader_BIN, test_loader_BIN = loaders_BIN
run_time_str = run_time_str if run_time_str is not None else datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
for seed in n_seeds:
for config in configs:
dir_name = f"{config.name}_{run_time_str}_{seed}"
store = Store(save_dir_artifacts, dir_name)
writer = store.tensorboard
if adv_acc:
metrics = {
'test_accuracy_True': float,
'test_accuracy_False': float,
'epoch': int}
else:
metrics = {
'test_accuracy_False': float,
'epoch': int}
store.add_table('result', metrics)
# std_linear_net = utils.Linear(Nfeatures=3*32*32, Nclasses=2).cuda()
model = get_new_model(config.robust_model)
if not config.fine_tune:
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Linear(512, len(targets)).cuda()
model.train()
train_model(model,
train_loader_BIN,
test_loader_BIN,
train_loop,
eval_loop,
step_size=config.step_size,
epochs=n_epochs,
normalization=normalization_function_BIN,
store=store,
adv_train=config.adv_train,
log_iterations=10,
adv_acc=adv_acc
)
if save_models:
path = f"{save_dir_models}/{dir_name}.pt"
torch.save(model.state_dict(), path)
def train_model(args,
model,
optim,
train_dl: DataLoader,
valid_dl: DataLoader,
store: Store,
device,
attack=None,
ratio: int = 0):
"""
Generic training routine, which is flexible to allow both standard and adversarial training.
"""
start_time = time.time()
# Initial setup
store.add_table(consts.LOGS_TABLE, consts.LOGS_SCHEMA)
store.add_table(consts.ARGS_TABLE, consts.ARGS_SCHEMA)
args_info = {
'epochs': args['--epochs'],
'batch_size': args['--batch-size'],
'model': 'mnist'
}
# store[consts.ARGS_TABLE].append_row(args_info)
model.to(device)
for epoch in range(args['--epochs']):
# Train for one epoch
train_acc, train_loss = _internal_loop(args, True, model, optim,
train_dl, epoch, store, device)
# Evaluate on validation
with torch.no_grad():
valid_acc, valid_loss = _internal_loop(args, False, model, None,
valid_dl, epoch, store,
device)
# Log
log_info = {
'epoch': epoch,
'train_loss': train_loss,
'valid_loss': valid_loss,
'train_error_rate': 1 - train_acc,
'valid_error_rate': 1 - valid_acc,
'valid_adv_error_rate': -1,
'time': time.time() - start_time
}
# store[consts.LOGS_TABLE].append_row(log_info)
return model
def initialize_cox_store(cox_dir='cox') -> Store:
store = Store(cox_dir)
store.add_table(
'experiments',
{
'k': int,
'random_state': int,
'Train AUC': float,
'Validation AUC': float,
'Test AUC': float,
'Test MCC': float,
'start_time': str,
# 'runtime(sec)': float,
'classifier': str,
'classifier_full': str
})
return store
def test_L1Ball(algorithm, step_size):
# Setup
constraint = constopt.constraints.L1Ball(alpha)
assert (constraint.prox(w) == w).all()
w_t = Variable(torch.zeros_like(w), requires_grad=True)
optimizer = algorithm([w_t], constraint)
criterion = torch.nn.MSELoss(reduction='mean')
# Logging
store = Store(OUT_DIR)
store.add_table('metadata', {'algorithm': str, 'step-size': float})
store['metadata'].append_row({
'algorithm': optimizer.name,
'step-size': step_size
})
store.add_table(optimizer.name, {
'func_val': float,
'FW gap': float,
'norm(w_t)': float
})
gap = torch.tensor(np.inf)
for ii in range(MAX_ITER):
optimizer.zero_grad()
loss = criterion(X.mv(w_t), y)
loss.backward()
# Compute gap
with torch.no_grad():
gap = constraint.fw_gap(w_t.grad, w_t)
optimizer.step(step_size)
store.log_table_and_tb(
optimizer.name, {
'func_val': loss.item(),
'FW gap': gap.item(),
'norm(w_t)': sum(abs(w_t)).item()
})
store[optimizer.name].flush_row()
store.close()
def test_L1Ball(algorithm, lr):
# Setup
constraint = chop.constraints.L1Ball(alpha)
assert (constraint.prox(w) == w).all()
w_t = Variable(torch.zeros_like(w), requires_grad=True)
optimizer = algorithm([w_t], constraint, lr=lr)
criterion = torch.nn.MSELoss(reduction='mean')
# Logging
store = Store(OUT_DIR)
store.add_table('metadata', {'algorithm': str, 'lr': float})
store['metadata'].append_row({'algorithm': optimizer.name, 'lr': lr})
store.add_table(optimizer.name, {
'func_val': float,
'certificate': float,
'norm(w_t)': float
})
cert = torch.tensor(np.inf)
for ii in range(MAX_ITER):
optimizer.zero_grad()
loss = criterion(X.mv(w_t), y)
loss.backward()
optimizer.step()
cert = next(optimizer.certificate) # only one parameter here
store.log_table_and_tb(
optimizer.name, {
'func_val': loss.item(),
'certificate': cert.item(),
'norm(w_t)': sum(abs(w_t)).item()
})
store[optimizer.name].flush_row()
store.close()
def main(params):
for k, v in zip(params.keys(), params.values()):
assert v is not None, f"Value for {k} is None"
# #
# Setup logging
# #
metadata_schema = schema_from_dict(params)
base_directory = params['out_dir']
store = Store(base_directory)
# redirect stderr, stdout to file
"""
def make_err_redirector(stream_name):
tee = Tee(os.path.join(store.path, stream_name + '.txt'), stream_name)
return tee
stderr_tee = make_err_redirector('stderr')
stdout_tee = make_err_redirector('stdout')
"""
# Store the experiment path and the git commit for this experiment
metadata_schema.update({
'store_path':str,
'git_commit':str
})
repo = git.Repo(path=os.path.dirname(os.path.realpath(__file__)),
search_parent_directories=True)
metadata_table = store.add_table('metadata', metadata_schema)
metadata_table.update_row(params)
metadata_table.update_row({
'store_path':store.path,
'git_commit':repo.head.object.hexsha
})
metadata_table.flush_row()
# Table for checkpointing models and envs
if params['save_iters'] > 0:
store.add_table('checkpoints', {
'val_model':store.PYTORCH_STATE,
'policy_model':store.PYTORCH_STATE,
'envs':store.PICKLE,
'policy_opt': store.PYTORCH_STATE,
'val_opt': store.PYTORCH_STATE,
'iteration':int
})
# The trainer object is in charge of sampling trajectories and
# taking PPO/TRPO optimization steps
p = Trainer.agent_from_params(params, store=store)
if 'load_model' in params and params['load_model']:
print('Loading pretrained model', params['load_model'])
pretrained_models = torch.load(params['load_model'])
p.policy_model.load_state_dict(pretrained_models['policy_model'])
p.val_model.load_state_dict(pretrained_models['val_model'])
# Load optimizer states. Note that
# p.POLICY_ADAM.load_state_dict(pretrained_models['policy_opt'])
# p.val_opt.load_state_dict(pretrained_models['val_opt'])
# Restore environment parameters, like mean and std.
p.envs = pretrained_models['envs']
rewards = []
# Table for final results
final_table = store.add_table('final_results', {
'iteration':int,
'5_rewards':float,
'terminated_early':bool,
'val_model':store.PYTORCH_STATE,
'policy_model':store.PYTORCH_STATE,
'envs':store.PICKLE,
'policy_opt': store.PYTORCH_STATE,
'val_opt': store.PYTORCH_STATE,
'iteration':int
})
def finalize_table(iteration, terminated_early, rewards):
final_5_rewards = np.array(rewards)[-5:].mean()
final_table.append_row({
'iteration':iteration,
'5_rewards':final_5_rewards,
'terminated_early':terminated_early,
'iteration':iteration,
'val_model': p.val_model.state_dict(),
'policy_model': p.policy_model.state_dict(),
'policy_opt': p.POLICY_ADAM.state_dict(),
'val_opt': p.val_opt.state_dict(),
'envs':p.envs
})
# Try-except so that we save if the user interrupts the process
try:
for i in range(params['train_steps']):
print('Step %d' % (i,))
if params['save_iters'] > 0 and i % params['save_iters'] == 0:
store['checkpoints'].append_row({
'iteration':i,
'val_model': p.val_model.state_dict(),
'policy_model': p.policy_model.state_dict(),
'policy_opt': p.POLICY_ADAM.state_dict(),
'val_opt': p.val_opt.state_dict(),
'envs':p.envs
})
mean_reward = p.train_step()
rewards.append(mean_reward)
finalize_table(i, False, rewards)
except KeyboardInterrupt:
torch.save(p.val_model, 'saved_experts/%s-expert-vf' % (params['game'],))
torch.save(p.policy_model, 'saved_experts/%s-expert-pol' % (params['game'],))
finalize_table(i, True, rewards)
store.close()
## Code sample to go alongside Walkthrough #2 in README.md
OUT_DIR = '/tmp/cox_example/'
try:
shutil.rmtree(OUT_DIR)
except:
pass
os.mkdir(OUT_DIR)
if __name__ == "__main__":
for slope in range(5):
store = Store(OUT_DIR)
store.add_table('metadata', {'slope': int})
store.add_table('line_graphs', {'mx': int, 'mx^2': int})
store['metadata'].append_row({'slope': slope})
for x in range(100):
store.log_table_and_tb('line_graphs', {
'mx': slope * x,
'mx^2': slope * (x**2)
})
store['line_graphs'].flush_row()
store.close()
### Collection reading
print("Done experiments, printing results...")
reader = CollectionReader(OUT_DIR)
def f(x):
return (x - 2.03)**2 + 3
if __name__ == "__main__":
# Three parameters: initial guess for x, step size, tolerance
combos = itertools.product(np.linspace(-15, 15, 3),
np.linspace(1, 1e-5, 3), [1e-5])
for x, step, tol in combos:
store = Store(OUT_DIR)
store.add_table(
'metadata', {
'step_size': float,
'tolerance': float,
'initial_x': float,
'out_dir': str
})
store.add_table('result', {'final_x': float, 'final_opt': float})
store.add_table('running_log', {
'current_x': float,
'current_f': float
})
store['metadata'].append_row({
'step_size': step,
'tolerance': tol,
'initial_x': x,
def main(params):
for k, v in zip(params.keys(), params.values()):
assert v is not None, f"Value for {k} is None"
# #
# Setup logging
# #
metadata_schema = schema_from_dict(params)
base_directory = params['out_dir']
store = Store(base_directory)
# redirect stderr, stdout to file
"""
def make_err_redirector(stream_name):
tee = Tee(os.path.join(store.path, stream_name + '.txt'), stream_name)
return tee
stderr_tee = make_err_redirector('stderr')
stdout_tee = make_err_redirector('stdout')
"""
# Store the experiment path and the git commit for this experiment
metadata_schema.update({
'store_path': str,
'git_commit': str
})
repo = git.Repo(path=os.path.dirname(os.path.realpath(__file__)),
search_parent_directories=True)
metadata_table = store.add_table('metadata', metadata_schema)
metadata_table.update_row(params)
metadata_table.update_row({
'store_path': store.path,
'git_commit': repo.head.object.hexsha
})
metadata_table.flush_row()
# Extra items in table when minimax training is enabled.
if params['mode'] == "adv_ppo" or params['mode'] == 'adv_trpo' or params['mode'] == 'adv_sa_ppo':
adversary_table_dict = {
'adversary_policy_model': store.PYTORCH_STATE,
'adversary_policy_opt': store.PYTORCH_STATE,
'adversary_val_model': store.PYTORCH_STATE,
'adversary_val_opt': store.PYTORCH_STATE,
}
else:
adversary_table_dict = {}
# Table for checkpointing models and envs
if params['save_iters'] > 0:
checkpoint_dict = {
'val_model': store.PYTORCH_STATE,
'policy_model': store.PYTORCH_STATE,
'envs': store.PICKLE,
'policy_opt': store.PYTORCH_STATE,
'val_opt': store.PYTORCH_STATE,
'iteration': int,
'5_rewards': float,
}
checkpoint_dict.update(adversary_table_dict)
store.add_table('checkpoints', checkpoint_dict)
# The trainer object is in charge of sampling trajectories and
# taking PPO/TRPO optimization steps
p = Trainer.agent_from_params(params, store=store)
if params['initial_std'] != 1.0:
p.policy_model.log_stdev.data[:] = np.log(params['initial_std'])
if 'load_model' in params and params['load_model']:
print('Loading pretrained model', params['load_model'])
pretrained_model = torch.load(params['load_model'])
if 'policy_model' in pretrained_model:
p.policy_model.load_state_dict(pretrained_model['policy_model'])
if params['deterministic']:
print('Policy runs in deterministic mode. Ignoring Gaussian noise.')
p.policy_model.log_stdev.data[:] = -100
else:
print('Policy runs in non deterministic mode with Gaussian noise.')
if 'val_model' in pretrained_model:
p.val_model.load_state_dict(pretrained_model['val_model'])
if 'policy_opt' in pretrained_model:
p.POLICY_ADAM.load_state_dict(pretrained_model['policy_opt'])
if 'val_opt' in pretrained_model:
p.val_opt.load_state_dict(pretrained_model['val_opt'])
# Load adversary models.
if 'no_load_adv_policy' in params and params['no_load_adv_policy']:
print('Skipping loading adversary models.')
else:
if 'adversary_policy_model' in pretrained_model and hasattr(p, 'adversary_policy_model'):
p.adversary_policy_model.load_state_dict(pretrained_model['adversary_policy_model'])
if 'adversary_val_model' in pretrained_model and hasattr(p, 'adversary_val_model'):
p.adversary_val_model.load_state_dict(pretrained_model['adversary_val_model'])
if 'adversary_policy_opt' in pretrained_model and hasattr(p, 'adversary_policy_opt'):
p.adversary_policy_opt.load_state_dict(pretrained_model['adversary_policy_opt'])
if 'adversary_val_opt' in pretrained_model and hasattr(p, 'adversary_val_opt'):
p.adversary_val_opt.load_state_dict(pretrained_model['adversary_val_opt'])
# Load optimizer states.
# p.POLICY_ADAM.load_state_dict(pretrained_models['policy_opt'])
# p.val_opt.load_state_dict(pretrained_models['val_opt'])
# Restore environment parameters, like mean and std.
if 'envs' in pretrained_model:
p.envs = pretrained_model['envs']
for e in p.envs:
e.setup_visualization(params['show_env'], params['save_frames'], params['save_frames_path'])
rewards = []
# Table for final results
final_dict = {
'iteration': int,
'5_rewards': float,
'terminated_early': bool,
'val_model': store.PYTORCH_STATE,
'policy_model': store.PYTORCH_STATE,
'envs': store.PICKLE,
'policy_opt': store.PYTORCH_STATE,
'val_opt': store.PYTORCH_STATE,
}
final_dict.update(adversary_table_dict)
final_table = store.add_table('final_results', final_dict)
def add_adversary_to_table(p, table_dict):
if params['mode'] == "adv_ppo" or params['mode'] == 'adv_trpo' or params['mode'] == 'adv_sa_ppo':
table_dict["adversary_policy_model"] = p.adversary_policy_model.state_dict()
table_dict["adversary_policy_opt"] = p.ADV_POLICY_ADAM.state_dict()
table_dict["adversary_val_model"] = p.adversary_val_model.state_dict()
table_dict["adversary_val_opt"] = p.adversary_val_opt.state_dict()
return table_dict
def finalize_table(iteration, terminated_early, rewards):
final_5_rewards = np.array(rewards)[-5:].mean()
final_dict = {
'iteration': iteration,
'5_rewards': final_5_rewards,
'terminated_early': terminated_early,
'val_model': p.val_model.state_dict(),
'policy_model': p.policy_model.state_dict(),
'policy_opt': p.POLICY_ADAM.state_dict(),
'val_opt': p.val_opt.state_dict(),
'envs': p.envs
}
final_dict = add_adversary_to_table(p, final_dict)
final_table.append_row(final_dict)
ret = 0
# Try-except so that we save if the user interrupts the process
try:
for i in range(params['train_steps']):
print('Step %d' % (i,))
if params['save_iters'] > 0 and i % params['save_iters'] == 0 and i != 0:
final_5_rewards = np.array(rewards)[-5:].mean()
print(f'Saving checkpoints to {store.path} with reward {final_5_rewards:.5g}')
checkpoint_dict = {
'iteration': i,
'val_model': p.val_model.state_dict(),
'policy_model': p.policy_model.state_dict(),
'policy_opt': p.POLICY_ADAM.state_dict(),
'val_opt': p.val_opt.state_dict(),
'envs': p.envs,
'5_rewards': final_5_rewards,
}
checkpoint_dict = add_adversary_to_table(p, checkpoint_dict)
store['checkpoints'].append_row(checkpoint_dict)
mean_reward = p.train_step()
rewards.append(mean_reward)
# For debugging and tuning, we can break in the middle.
if i == params['force_stop_step']:
print('Terminating early because --force-stop-step is set.')
raise KeyboardInterrupt
finalize_table(i, False, rewards)
except KeyboardInterrupt:
finalize_table(i, True, rewards)
ret = 1
except:
print("An error occurred during training:")
traceback.print_exc()
# Other errors, make sure to finalize the cox store before exiting.
finalize_table(i, True, rewards)
ret = -1
print(f'Models saved to {store.path}')
store.close()
return ret