def main(epochs, save_every, _log, _run):
run_dir = get_run_dir(_run)
os.makedirs(os.path.join(run_dir, "model_statedicts"))
dset, train_loader, val_loader, test = make_dataloaders()
model = make_model()
optimizer = make_optimizer(model)
lossfn = nn.CrossEntropyLoss()
for epoch in range(epochs + 1):
# train
t = tqdm(train_loader)
t.set_description(desc=f"Epoch: {epoch}")
for x, y in t:
y_pred = model(x)
loss = lossfn(y_pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
t.set_postfix(loss=loss.item())
ex.log_scalar("train_loss", loss.item(), epoch)
# validate
with torch.no_grad():
model.eval()
total_loss = 0
total_examples = 0
for x, y in val_loader:
y_pred = model(x)
loss = lossfn(y_pred, y)
total_loss += loss.item() * len(y)
total_examples += len(y)
val_loss = total_loss / total_examples
_log.info(f"val_loss={val_loss:.6f}")
ex.log_scalar("val_loss", val_loss, epoch)
model.train()
# save
if epoch % save_every == 0:
name = f"{epoch:03d}model_val{val_loss:.4f}_statedict.pickle"
path = os.path.join(run_dir, "model_statedicts", name)
torch.save(model.state_dict(), path)
ex.add_artifact(path, name)
def main(_run):
dset, train, val, test = make_dataloaders()
model = make_model()
optimizer = make_optimizer(model)
callback = make_scheduler_callback(optimizer)
st.loop(
**{**_run.config,
**dict(_run=_run,
model=model,
optimizer=optimizer,
save_dir=SAVE_DIR,
trainOnBatch=train_on_batch,
train_loader=train,
val_loader=val,
callback=callback,
callback_metric_names=['val_loss', 'val_acc', 'learning_rate'],
batch_metric_names=['loss', 'acc', 'nfef', 'nfeb'],
updaters=[averager]*4)})
def main(_run,
attack,
epsilon,
pgd_step_size,
pgd_num_steps,
pgd_random_start,):
dset, train, val, test = make_dataloaders()
model = make_model()
optimizer = make_optimizer(model)
callback = make_scheduler_callback(optimizer)
if attack == 'pgd':
attack_fn = partial(ATTACKS[attack],
epsilon=epsilon,
step_size=pgd_step_size,
num_steps=pgd_num_steps,
random_start=pgd_random_start)
else:
attack_fn = partial(ATTACKS[attack], epsilon=epsilon)
adv_train = AdversarialLoader(model, train, attack_fn)
final_train = CombineDataloaders(train, adv_train)
st.loop(
**{**_run.config,
**dict(_run=_run,
model=model,
optimizer=optimizer,
save_dir=SAVE_DIR,
trainOnBatch=train_on_batch,
train_loader=final_train,
val_loader=val,
callback=callback,
callback_metric_names=['val_loss', 'val_acc', 'learning_rate'],
batch_metric_names=['loss', 'acc', 'nfef', 'nfeb'],
updaters=[averager]*4)})