def main(cfg):
dat = HubmapDataset(cfg["data_dir"], cfg["out_dir"])
db = zarr.open_group(store=zarr.DirectoryStore(cfg["zarr_db_dir"]),
mode="r")
# results dir
results_dir = dat.path.out / "models" / cfg["version"]
if cfg["version"] == "debug":
os.makedirs(results_dir, exist_ok=True)
else:
try:
os.makedirs(results_dir, exist_ok=False)
except:
raise Exception(f"Version {cfg['version']} exists!")
with open(results_dir / "config.json", "w") as f:
json.dump(cfg, f, sort_keys=True, indent=4)
# tiles
tile_dct = dat.pkl_load(dat.path.out / "tiles" / cfg['tiles_version'] /
"tile_dct.pkl")
train_df = tile_dct["train_df"]
valid_dct = tile_dct["valid_dct"]
valid_id = list(valid_dct.keys())
train_valid_df = tile_dct["train_valid_df"]
# -- debug (cutting data)
if cfg["debug"]:
train_df = train_df.head(cfg["debug"])
valid_dct = {
id_: valid_dct[id_].head(cfg["debug"])
for id_ in valid_id
}
train_valid_df = train_valid_df.head(cfg["debug"])
# PyTorch Datasets initialization
augmentation = get_augmentation(cfg["augmentation_version"])
dataset = {
id_: ValidDataset(db, valid_dct[id_], augmentation["valid"])
for id_ in valid_id
}
dataset["train"] = TrainDataset(db, train_df, augmentation["train"])
dataset["train_valid"] = ValidDataset(db, train_valid_df,
augmentation["valid"])
# PyTorch DataLoaders initialization
dataloader = {
id_: DataLoader(dataset[id_],
batch_size=cfg["batch_size"],
shuffle=False,
num_workers=cfg["n_jobs"])
for id_ in valid_id
}
dataloader["train"] = DataLoader(dataset["train"],
batch_size=cfg["batch_size"],
shuffle=True,
num_workers=cfg["n_jobs"])
dataloader["train_valid"] = DataLoader(dataset["train_valid"],
batch_size=cfg["batch_size"],
shuffle=True,
num_workers=cfg["n_jobs"])
# model
model = get_model(cfg["model_version"])
if cfg["model_weights"]:
model.load_state_dict(torch.load(cfg["model_weights"]))
model.to(cfg["device"])
# optimizer
optimizer = get_optimizer(cfg["optimizer_version"], model.parameters())
if cfg["optimizer_weights"]:
optimizer.load_state_dict(torch.load(cfg["optimizer_weights"]))
# scheduler
if cfg["scheduler_version"]:
scheduler = get_scheduler(cfg["scheduler_version"], optimizer)
# loss function
criterion = get_criterion(cfg["criterion_version"])
# EPOCHS
monitor = []
loss_min = None
dice_max = None
epochs_without_improvement = 0
for epoch in range(cfg["epoch_num"]):
print(f"Epoch-{epoch}")
monitor_epoch = {}
# train
train_loss_raw = train(model,
dataloader["train"],
criterion,
optimizer,
device=cfg["device"],
verbose="train-train")
monitor_epoch["train_loss_raw"] = train_loss_raw
# train loss
_, train_loss = valid(model,
dataloader["train_valid"],
criterion,
dice_ths=False,
device=cfg["device"],
verbose="train-valid")
monitor_epoch["train_loss"] = train_loss
# scheduler
if cfg["scheduler_version"]:
scheduler.step(train_loss)
# validation
valid_loss = 0
valid_dice = 0
for id_ in valid_id:
valid_dice_id, valid_loss_id = valid(model,
dataloader[id_],
criterion,
dice_ths=cfg["dice_ths"],
device=cfg["device"],
verbose=id_)
valid_dice += valid_dice_id
valid_loss += valid_loss_id
valid_loss /= len(valid_id)
valid_dice /= len(valid_id)
monitor_epoch["valid_loss"] = valid_loss
monitor_epoch["valid_dice"] = valid_dice
print(
f"Train-loss: {train_loss:.5} \t Valid-loss: {valid_loss:.5} \t Valid-dice: {valid_dice:.3}"
)
print("-" * 70)
# saving progress info
monitor.append(monitor_epoch)
dat.pkl_dump(pd.DataFrame(monitor), results_dir / "monitor.pkl")
# saving weights - max DICE
if dice_max is None:
dice_max = valid_dice
if valid_dice >= dice_max:
dice_max = valid_dice
torch.save(model.state_dict(), results_dir / "model_best_dice.pt")
# saving weights - min LOSS
if loss_min is None:
loss_min = valid_loss
# loss improvement
if valid_loss <= loss_min:
loss_min = valid_loss
epochs_without_improvement = 0
# save model
torch.save(model.state_dict(), results_dir / "model_best_loss.pt")
torch.save(optimizer.state_dict(), results_dir / "optimizer.pt")
# -- no loss improvement
else:
epochs_without_improvement += 1
# early stopping
if epochs_without_improvement >= cfg["early_stopping"]:
print("EARLY STOPPING!")
break
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config', '-c', type=str, required=True)
parser.add_argument('--device', '-d', default='cuda')
parser.add_argument('--batch_size', '-b', type=int, default=256)
parser.add_argument('--epochs', '-e', type=int, default=400)
parser.add_argument('--stopping_param', type=int, default=50)
parser.add_argument('--save_iter', '-s', type=int, default=10)
parser.add_argument('--vis_iter', '-v', type=int, default=10)
parser.add_argument('--log_dir', '-l', type=str, default=None)
parser.add_argument('--seed', type=int, default=42)
parser.add_argument(
'--dataset',
'-D',
type=str,
default='mnist',
choices=['mnist', 'cifar10', 'cifar100', 'clothing1m', 'imagenet'])
parser.add_argument('--data_augmentation',
'-A',
action='store_true',
dest='data_augmentation')
parser.set_defaults(data_augmentation=False)
parser.add_argument('--num_train_examples', type=int, default=None)
parser.add_argument('--label_noise_level', '-n', type=float, default=0.0)
parser.add_argument('--label_noise_type',
type=str,
default='error',
choices=['error', 'cifar10_custom'])
parser.add_argument('--transform_function',
type=str,
default=None,
choices=[None, 'remove_random_chunks'])
parser.add_argument('--clean_validation',
dest='clean_validation',
action='store_true')
parser.set_defaults(clean_validation=False)
parser.add_argument('--remove_prob', type=float, default=0.5)
parser.add_argument('--model_class',
'-m',
type=str,
default='StandardClassifier')
parser.add_argument(
'--loss_function',
type=str,
default='ce',
choices=['ce', 'mse', 'mae', 'gce', 'dmi', 'fw', 'none'])
parser.add_argument('--loss_function_param', type=float, default=1.0)
parser.add_argument('--load_from', type=str, default=None)
parser.add_argument('--grad_weight_decay', '-L', type=float, default=0.0)
parser.add_argument('--grad_l1_penalty', '-S', type=float, default=0.0)
parser.add_argument('--lamb', type=float, default=1.0)
parser.add_argument('--pretrained_arg', '-r', type=str, default=None)
parser.add_argument('--sample_from_q',
action='store_true',
dest='sample_from_q')
parser.set_defaults(sample_from_q=False)
parser.add_argument('--q_dist',
type=str,
default='Gaussian',
choices=['Gaussian', 'Laplace', 'dot'])
parser.add_argument('--no-detach', dest='detach', action='store_false')
parser.set_defaults(detach=True)
parser.add_argument('--warm_up',
type=int,
default=0,
help='Number of epochs to skip before '
'starting to train using predicted gradients')
parser.add_argument('--weight_decay', type=float, default=0.0)
parser.add_argument(
'--add_noise',
action='store_true',
dest='add_noise',
help='add noise to the gradients of a standard classifier.')
parser.set_defaults(add_noise=False)
parser.add_argument('--noise_type',
type=str,
default='Gaussian',
choices=['Gaussian', 'Laplace'])
parser.add_argument('--noise_std', type=float, default=0.0)
parser.add_argument('--lr', type=float, default=1e-3, help='Learning rate')
args = parser.parse_args()
print(args)
# Load data
train_loader, val_loader, test_loader = datasets.load_data_from_arguments(
args)
# Options
optimization_args = {
'optimizer': {
'name': 'adam',
'lr': args.lr,
'weight_decay': args.weight_decay
}
}
# optimization_args = {
# 'optimizer': {
# 'name': 'sgd',
# 'lr': 1e-3,
# },
# 'scheduler': {
# 'step_size': 15,
# 'gamma': 1.25
# }
# }
with open(args.config, 'r') as f:
architecture_args = json.load(f)
model_class = getattr(methods, args.model_class)
model = model_class(input_shape=train_loader.dataset[0][0].shape,
architecture_args=architecture_args,
pretrained_arg=args.pretrained_arg,
device=args.device,
grad_weight_decay=args.grad_weight_decay,
grad_l1_penalty=args.grad_l1_penalty,
lamb=args.lamb,
sample_from_q=args.sample_from_q,
q_dist=args.q_dist,
load_from=args.load_from,
loss_function=args.loss_function,
loss_function_param=args.loss_function_param,
add_noise=args.add_noise,
noise_type=args.noise_type,
noise_std=args.noise_std,
detach=args.detach,
warm_up=args.warm_up)
metrics_list = []
if args.dataset == 'imagenet':
metrics_list.append(metrics.TopKAccuracy(k=5, output_key='pred'))
training.train(model=model,
train_loader=train_loader,
val_loader=val_loader,
epochs=args.epochs,
save_iter=args.save_iter,
vis_iter=args.vis_iter,
optimization_args=optimization_args,
log_dir=args.log_dir,
args_to_log=args,
stopping_param=args.stopping_param,
metrics=metrics_list)
# if training finishes successfully, compute the test score
print("Testing the best validation model...")
model = utils.load(os.path.join(args.log_dir, 'checkpoints',
'best_val.mdl'),
device=args.device)
pred = utils.apply_on_dataset(model,
test_loader.dataset,
batch_size=args.batch_size,
output_keys_regexp='pred',
description='Testing')['pred']
labels = [p[1] for p in test_loader.dataset]
labels = torch.tensor(labels, dtype=torch.long)
labels = utils.to_cpu(labels)
with open(os.path.join(args.log_dir, 'test_predictions.pkl'), 'wb') as f:
pickle.dump({'pred': pred, 'labels': labels}, f)
accuracy = torch.mean((pred.argmax(dim=1) == labels).float())
with open(os.path.join(args.log_dir, 'test_accuracy.txt'), 'w') as f:
f.write("{}\n".format(accuracy))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", "-c", type=str, required=True)
parser.add_argument("--device", "-d", default="cuda")
parser.add_argument("--batch_size", "-b", type=int, default=256)
parser.add_argument("--epochs", "-e", type=int, default=400)
parser.add_argument("--save_iter", "-s", type=int, default=10)
parser.add_argument("--vis_iter", "-v", type=int, default=10)
parser.add_argument("--log_dir", "-l", type=str, default=None)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--dataset",
"-D",
type=str,
default="mnist",
choices=["mnist", "cifar10"])
parser.add_argument("--data_augmentation",
"-A",
action="store_true",
dest="data_augmentation")
parser.set_defaults(data_augmentation=False)
parser.add_argument("--num_train_examples", type=int, default=None)
parser.add_argument("--label_noise_level", "-n", type=float, default=0.0)
parser.add_argument(
"--label_noise_type",
type=str,
default="flip",
choices=["flip", "error", "cifar10_custom"],
)
parser.add_argument(
"--transform_function",
type=str,
default=None,
choices=[None, "remove_random_chunks"],
)
parser.add_argument("--clean_validation",
dest="clean_validation",
action="store_true")
parser.set_defaults(clean_validation=False)
parser.add_argument("--remove_prob", type=float, default=0.5)
args = parser.parse_args()
print(args)
# Load data
train_loader, val_loader, test_loader = datasets.load_data_from_arguments(
args)
# Options
optimization_args = {"optimizer": {"name": "adam", "lr": 1e-3}}
with open(args.config, "r") as f:
architecture_args = json.load(f)
revert_normalization = lambda x: datasets.revert_normalization(
x, train_loader.dataset)
model = VAE(
input_shape=train_loader.dataset[0][0].shape,
architecture_args=architecture_args,
device=args.device,
revert_normalization=revert_normalization,
)
training.train(
model=model,
train_loader=train_loader,
val_loader=val_loader,
epochs=args.epochs,
save_iter=args.save_iter,
vis_iter=args.vis_iter,
optimization_args=optimization_args,
log_dir=args.log_dir,
)
def main(cfg):
dat = HubmapDataset(cfg["data_dir"], cfg["out_dir"])
db = zarr.open_group(store=zarr.DirectoryStore(cfg["zarr_db_dir"]),
mode="r")
# results dir
results_dir = dat.path.out / "seed_selection" / cfg["version"]
if cfg["version"] == "debug":
os.makedirs(results_dir, exist_ok=True)
else:
try:
os.makedirs(results_dir, exist_ok=False)
except:
raise Exception(f"Version {cfg['version']} exists!")
# save config
dat.jsn_dump(cfg, results_dir / "config.json")
# tiles
tile_dct = dat.pkl_load(dat.path.out / "tiles" / cfg['tiles_version'] /
"tile_dct.pkl")
# train/valid split
valid_dct = tile_dct["valid_dct"]
valid_id = list(valid_dct.keys())
train_df = tile_dct["train_df"]
# -- debug (cutting data)
if cfg["debug"]:
valid_dct = {
id_: valid_dct[id_].head(cfg["debug"])
for id_ in valid_id
}
train_df = train_df.head(cfg["debug"])
# PyTorch Datasets initialization
augmentation = get_augmentation(cfg["augmentation_version"])
dataset = {
id_: ValidDataset(db, valid_dct[id_], augmentation["valid"])
for id_ in valid_id
}
dataset["train"] = TrainDataset(db, train_df, augmentation["valid"])
# PyTorch DataLoaders initialization
dataloader = {
id_: DataLoader(dataset[id_],
batch_size=cfg["batch_size"],
shuffle=True,
num_workers=1)
for id_ in valid_id
}
dataloader["train"] = DataLoader(dataset["train"],
batch_size=cfg["batch_size"],
shuffle=True,
num_workers=1)
# seed search
time_start = time.time()
time_stop = False
seeds = []
valid_loss_min = None
while time_stop is not True:
# seed value
while True:
seed = random.randint(0, 100_000)
if seed not in seeds:
seeds.append(seed)
break
seed_everything(seed)
# model initialization
model = get_model(cfg["model_version"])
# loss-function
criterion = get_criterion(cfg["criterion_version"])
# training
optimizer = get_optimizer(cfg["optimizer_version"], model.parameters())
for epoch in range(cfg["epoch_num"]):
train(model, dataloader["train"], criterion, optimizer)
# validation
valid_loss = 0
for id_ in valid_id:
_, valid_loss_id = valid(model,
dataloader[id_],
criterion,
dice_ths=cfg["dice_ths"],
device=cfg["device"],
verbose=id_)
valid_loss += valid_loss_id
valid_loss /= len(valid_id)
print(f"Seed-{seed} | Valid-loss: {valid_loss:.5}")
# saving progress
if valid_loss_min is None:
valid_loss_min = valid_loss
if valid_loss <= valid_loss_min:
valid_loss_min = valid_loss
torch.save(model.state_dict(), results_dir / "model.pt")
with open(results_dir / "best_seed.txt", "w") as f:
f.write(str(seed))
# stop by time
if (time.time() - time_start) / 60 >= cfg["max_time"]:
time_stop = True
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", "-c", type=str, required=True)
parser.add_argument("--device", "-d", default="cuda")
parser.add_argument("--batch_size", "-b", type=int, default=128)
parser.add_argument("--epochs", "-e", type=int, default=4000)
parser.add_argument("--stopping_param", type=int, default=2**30)
parser.add_argument("--save_iter", "-s", type=int, default=100)
parser.add_argument("--vis_iter", "-v", type=int, default=10)
parser.add_argument("--log_dir", "-l", type=str, default=None)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument(
"--dataset",
"-D",
type=str,
default="cifar10",
choices=["mnist", "cifar10", "cifar100", "clothing1m", "imagenet"],
)
parser.add_argument("--data_augmentation",
"-A",
action="store_true",
dest="data_augmentation")
parser.set_defaults(data_augmentation=False)
parser.add_argument("--num_train_examples", type=int, default=None)
parser.add_argument("--label_noise_level", "-n", type=float, default=0.0)
parser.add_argument(
"--label_noise_type",
type=str,
default="error",
choices=["error", "cifar10_custom"],
)
parser.add_argument(
"--transform_function",
type=str,
default=None,
choices=[None, "remove_random_chunks"],
)
parser.add_argument("--clean_validation",
dest="clean_validation",
action="store_true")
parser.set_defaults(clean_validation=False)
parser.add_argument("--remove_prob", type=float, default=0.5)
parser.add_argument("--model_class",
"-m",
type=str,
default="StandardClassifier")
parser.add_argument("--load_from", type=str, default=None)
parser.add_argument("--grad_weight_decay", "-L", type=float, default=0.0)
parser.add_argument("--lamb", type=float, default=1.0)
parser.add_argument("--pretrained_arg", "-r", type=str, default=None)
parser.add_argument("--sample_from_q",
action="store_true",
dest="sample_from_q")
parser.set_defaults(sample_from_q=False)
parser.add_argument("--q_dist",
type=str,
default="Gaussian",
choices=["Gaussian", "Laplace", "dot"])
parser.add_argument("--weight_decay", type=float, default=0.0)
parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
parser.add_argument(
"--k",
"-k",
type=int,
required=False,
default=10,
help="width parameter of ResNet18-k",
)
parser.add_argument("--exclude_percent", type=float, default=0.0)
args = parser.parse_args()
print(args)
# Load data
train_loader, val_loader, test_loader = datasets.load_data_from_arguments(
args)
# Options
optimization_args = {
"optimizer": {
"name": "adam",
"lr": args.lr,
"weight_decay": args.weight_decay
}
}
with open(args.config, "r") as f:
architecture_args = json.load(f)
# set the width parameter k
if ("classifier" in architecture_args
and architecture_args["classifier"].get(
"net", "").find("double-descent") != -1):
architecture_args["classifier"]["k"] = args.k
if ("q-network" in architecture_args
and architecture_args["classifier"].get(
"net", "").find("double-descent") != -1):
architecture_args["q-network"]["k"] = args.k
model_class = getattr(methods, args.model_class)
model = model_class(
input_shape=train_loader.dataset[0][0].shape,
architecture_args=architecture_args,
pretrained_arg=args.pretrained_arg,
device=args.device,
grad_weight_decay=args.grad_weight_decay,
lamb=args.lamb,
sample_from_q=args.sample_from_q,
q_dist=args.q_dist,
load_from=args.load_from,
loss_function="ce",
)
training.train(
model=model,
train_loader=train_loader,
val_loader=val_loader,
epochs=args.epochs,
save_iter=args.save_iter,
vis_iter=args.vis_iter,
optimization_args=optimization_args,
log_dir=args.log_dir,
args_to_log=args,
stopping_param=args.stopping_param,
)
# test the last model and best model
models_to_test = [
{
"name": "best",
"file": "best_val.mdl"
},
{
"name": "final",
"file": "final.mdl"
},
]
for spec in models_to_test:
print("Testing the {} model...".format(spec["name"]))
model = utils.load(os.path.join(args.log_dir, "checkpoints",
spec["file"]),
device=args.device)
pred = utils.apply_on_dataset(
model,
test_loader.dataset,
batch_size=args.batch_size,
output_keys_regexp="pred",
description="Testing",
)["pred"]
labels = [p[1] for p in test_loader.dataset]
labels = torch.tensor(labels, dtype=torch.long)
labels = utils.to_cpu(labels)
with open(
os.path.join(args.log_dir,
"{}_test_predictions.pkl".format(spec["name"])),
"wb",
) as f:
pickle.dump({"pred": pred, "labels": labels}, f)
accuracy = torch.mean((pred.argmax(dim=1) == labels).float())
with open(
os.path.join(args.log_dir,
"{}_test_accuracy.txt".format(spec["name"])),
"w") as f:
f.write("{}\n".format(accuracy))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", "-c", type=str, default=None)
parser.add_argument("--device", "-d", default="cuda")
parser.add_argument("--batch_size", "-b", type=int, default=256)
parser.add_argument("--epochs", "-e", type=int, default=400)
parser.add_argument("--stopping_param", type=int, default=50)
parser.add_argument("--save_iter", "-s", type=int, default=10)
parser.add_argument("--vis_iter", "-v", type=int, default=10)
parser.add_argument("--log_dir", "-l", type=str, default=None)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument(
"--dataset",
"-D",
type=str,
default="mnist",
choices=[
"mnist", "cifar10", "cifar100", "clothing1m", "imagenet", "cover"
],
)
# parser.add_argument("--image-root", metavar="DIR", help="path to images")
# parser.add_argument("--label", metavar="DIR", help="path to label file")
parser.add_argument("--data_augmentation",
"-A",
action="store_true",
dest="data_augmentation")
parser.set_defaults(data_augmentation=False)
parser.add_argument("--num_train_examples", type=int, default=None)
parser.add_argument("--label_noise_level", "-n", type=float, default=0.0)
parser.add_argument(
"--label_noise_type",
type=str,
default="error",
choices=["error", "cifar10_custom"],
)
parser.add_argument(
"--transform_function",
type=str,
default=None,
choices=[None, "remove_random_chunks"],
)
parser.add_argument("--clean_validation",
dest="clean_validation",
action="store_true")
parser.set_defaults(clean_validation=False)
parser.add_argument("--remove_prob", type=float, default=0.5)
parser.add_argument("--model_class",
"-m",
type=str,
default="StandardClassifier")
parser.add_argument(
"--loss_function",
type=str,
default="ce",
choices=["ce", "mse", "mae", "gce", "dmi", "fw", "none"],
)
parser.add_argument("--loss_function_param", type=float, default=1.0)
parser.add_argument("--load_from", type=str, default=None)
parser.add_argument("--grad_weight_decay", "-L", type=float, default=0.0)
parser.add_argument("--grad_l1_penalty", "-S", type=float, default=0.0)
parser.add_argument("--lamb", type=float, default=1.0)
parser.add_argument("--pretrained_arg", "-r", type=str, default=None)
parser.add_argument("--sample_from_q",
action="store_true",
dest="sample_from_q")
parser.set_defaults(sample_from_q=False)
parser.add_argument("--q_dist",
type=str,
default="Gaussian",
choices=["Gaussian", "Laplace", "dot"])
parser.add_argument("--no-detach", dest="detach", action="store_false")
parser.set_defaults(detach=True)
parser.add_argument(
"--warm_up",
type=int,
default=0,
help="Number of epochs to skip before "
"starting to train using predicted gradients",
)
parser.add_argument("--weight_decay", type=float, default=0.0)
parser.add_argument(
"--add_noise",
action="store_true",
dest="add_noise",
help="add noise to the gradients of a standard classifier.",
)
parser.set_defaults(add_noise=False)
parser.add_argument("--noise_type",
type=str,
default="Gaussian",
choices=["Gaussian", "Laplace"])
parser.add_argument("--noise_std", type=float, default=0.0)
parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate")
args = parser.parse_args()
print(args)
# Load data
args.image_root = [
"/data/chenlong.1024/5198852-tiktok-1w_images",
"/data/chenlong.1024/5205046-tiktok-10w_images",
"/data/chenlong.1024/5599074-tiktok-impr_cnt20_images",
"/data/chenlong.1024/5600297-tiktok-impr_cnt10_images",
]
args.label = [
"/data/chenlong.1024/5198852-tiktok-1w.csv",
"/data/chenlong.1024/5205046-tiktok-10w.csv",
"/data/chenlong.1024/5599074-tiktok-impr_cnt20.csv",
"/data/chenlong.1024/5600297-tiktok-impr_cnt10.csv",
]
train_loader, val_loader, test_loader = datasets.load_data_from_arguments(
args)
# Options
optimization_args = {
"optimizer": {
"name": "adam",
"lr": args.lr,
"weight_decay": args.weight_decay
},
"scheduler": {
"step_size": 20,
"gamma": 0.3
},
}
# optimization_args = {
# 'optimizer': {
# 'name': 'sgd',
# 'lr': 1e-3,
# },
# 'scheduler': {
# 'step_size': 15,
# 'gamma': 1.25
# }
# }
model_class = getattr(methods, args.model_class)
if "CoverModel" in args.model_class:
model = model_class(
num_classes=2,
pretrained=True,
device=args.device,
grad_weight_decay=args.grad_weight_decay,
grad_l1_penalty=args.grad_l1_penalty,
lamb=args.lamb,
sample_from_q=args.sample_from_q,
q_dist=args.q_dist,
load_from=args.load_from,
loss_function=args.loss_function,
loss_function_param=args.loss_function_param,
add_noise=args.add_noise,
noise_type=args.noise_type,
noise_std=args.noise_std,
detach=args.detach,
warm_up=args.warm_up,
)
else:
with open(args.config, "r") as f:
architecture_args = json.load(f)
model = model_class(
input_shape=train_loader.dataset[0][0].shape,
architecture_args=architecture_args,
pretrained_arg=args.pretrained_arg,
device=args.device,
grad_weight_decay=args.grad_weight_decay,
grad_l1_penalty=args.grad_l1_penalty,
lamb=args.lamb,
sample_from_q=args.sample_from_q,
q_dist=args.q_dist,
load_from=args.load_from,
loss_function=args.loss_function,
loss_function_param=args.loss_function_param,
add_noise=args.add_noise,
noise_type=args.noise_type,
noise_std=args.noise_std,
detach=args.detach,
warm_up=args.warm_up,
)
metrics_list = []
if args.dataset == "imagenet":
metrics_list.append(metrics.TopKAccuracy(k=5, output_key="pred"))
training.train(
model=model,
train_loader=train_loader,
val_loader=val_loader,
epochs=args.epochs,
save_iter=args.save_iter,
vis_iter=args.vis_iter,
optimization_args=optimization_args,
log_dir=args.log_dir,
args_to_log=args,
stopping_param=args.stopping_param,
metrics=metrics_list,
)
# if training finishes successfully, compute the test score
print("Testing the best validation model...")
model = utils.load(os.path.join(args.log_dir, "checkpoints",
"best_val.mdl"),
device=args.device)
pred = utils.apply_on_dataset(
model,
test_loader.dataset,
batch_size=args.batch_size,
output_keys_regexp="pred",
description="Testing",
)["pred"]
labels = [p[1] for p in test_loader.dataset]
labels = torch.tensor(labels, dtype=torch.long)
labels = utils.to_cpu(labels)
with open(os.path.join(args.log_dir, "test_predictions.pkl"), "wb") as f:
pickle.dump({"pred": pred, "labels": labels}, f)
accuracy = torch.mean((pred.argmax(dim=1) == labels).float())
with open(os.path.join(args.log_dir, "test_accuracy.txt"), "w") as f:
f.write("{}\n".format(accuracy))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config', '-c', type=str, required=True)
parser.add_argument('--device', '-d', default='cuda')
parser.add_argument('--batch_size', '-b', type=int, default=128)
parser.add_argument('--epochs', '-e', type=int, default=4000)
parser.add_argument('--stopping_param', type=int, default=2**30)
parser.add_argument('--save_iter', '-s', type=int, default=100)
parser.add_argument('--vis_iter', '-v', type=int, default=10)
parser.add_argument('--log_dir', '-l', type=str, default=None)
parser.add_argument('--seed', type=int, default=42)
parser.add_argument(
'--dataset',
'-D',
type=str,
default='cifar10',
choices=['mnist', 'cifar10', 'cifar100', 'clothing1m', 'imagenet'])
parser.add_argument('--data_augmentation',
'-A',
action='store_true',
dest='data_augmentation')
parser.set_defaults(data_augmentation=False)
parser.add_argument('--num_train_examples', type=int, default=None)
parser.add_argument('--label_noise_level', '-n', type=float, default=0.0)
parser.add_argument('--label_noise_type',
type=str,
default='error',
choices=['error', 'cifar10_custom'])
parser.add_argument('--transform_function',
type=str,
default=None,
choices=[None, 'remove_random_chunks'])
parser.add_argument('--clean_validation',
dest='clean_validation',
action='store_true')
parser.set_defaults(clean_validation=False)
parser.add_argument('--remove_prob', type=float, default=0.5)
parser.add_argument('--model_class',
'-m',
type=str,
default='StandardClassifier')
parser.add_argument('--load_from', type=str, default=None)
parser.add_argument('--grad_weight_decay', '-L', type=float, default=0.0)
parser.add_argument('--lamb', type=float, default=1.0)
parser.add_argument('--pretrained_arg', '-r', type=str, default=None)
parser.add_argument('--sample_from_q',
action='store_true',
dest='sample_from_q')
parser.set_defaults(sample_from_q=False)
parser.add_argument('--q_dist',
type=str,
default='Gaussian',
choices=['Gaussian', 'Laplace', 'dot'])
parser.add_argument('--weight_decay', type=float, default=0.0)
parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
parser.add_argument('--k',
'-k',
type=int,
required=True,
default=10,
help='width parameter of ResNet18-k')
args = parser.parse_args()
print(args)
# Load data
train_loader, val_loader, test_loader = datasets.load_data_from_arguments(
args)
# Options
optimization_args = {
'optimizer': {
'name': 'adam',
'lr': args.lr,
'weight_decay': args.weight_decay
}
}
with open(args.config, 'r') as f:
architecture_args = json.load(f)
# set the width parameter k
if ('classifier' in architecture_args
and architecture_args['classifier'].get(
'net', '') == 'double-descent-cifar10-resnet18'):
architecture_args['classifier']['k'] = args.k
if ('q-network' in architecture_args
and architecture_args['q-network'].get(
'net', '') == 'double-descent-cifar10-resnet18'):
architecture_args['q-network']['k'] = args.k
model_class = getattr(methods, args.model_class)
model = model_class(input_shape=train_loader.dataset[0][0].shape,
architecture_args=architecture_args,
pretrained_arg=args.pretrained_arg,
device=args.device,
grad_weight_decay=args.grad_weight_decay,
lamb=args.lamb,
sample_from_q=args.sample_from_q,
q_dist=args.q_dist,
load_from=args.load_from,
loss_function='ce')
training.train(model=model,
train_loader=train_loader,
val_loader=val_loader,
epochs=args.epochs,
save_iter=args.save_iter,
vis_iter=args.vis_iter,
optimization_args=optimization_args,
log_dir=args.log_dir,
args_to_log=args,
stopping_param=args.stopping_param)
# test the last model and best model
models_to_test = [{
'name': 'best',
'file': 'best_val.mdl'
}, {
'name': 'final',
'file': 'final.mdl'
}]
for spec in models_to_test:
print("Testing the {} model...".format(spec['name']))
model = utils.load(os.path.join(args.log_dir, 'checkpoints',
spec['file']),
device=args.device)
pred = utils.apply_on_dataset(model,
test_loader.dataset,
batch_size=args.batch_size,
output_keys_regexp='pred',
description='Testing')['pred']
labels = [p[1] for p in test_loader.dataset]
labels = torch.tensor(labels, dtype=torch.long)
labels = utils.to_cpu(labels)
with open(
os.path.join(args.log_dir,
'{}_test_predictions.pkl'.format(spec['name'])),
'wb') as f:
pickle.dump({'pred': pred, 'labels': labels}, f)
accuracy = torch.mean((pred.argmax(dim=1) == labels).float())
with open(
os.path.join(args.log_dir,
'{}_test_accuracy.txt'.format(spec['name'])),
'w') as f:
f.write("{}\n".format(accuracy))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config', '-c', type=str, required=True)
parser.add_argument('--device', '-d', default='cuda')
parser.add_argument('--batch_size', '-b', type=int, default=256)
parser.add_argument('--epochs', '-e', type=int, default=400)
parser.add_argument('--save_iter', '-s', type=int, default=10)
parser.add_argument('--vis_iter', '-v', type=int, default=10)
parser.add_argument('--log_dir', '-l', type=str, default=None)
parser.add_argument('--seed', type=int, default=42)
parser.add_argument('--dataset',
'-D',
type=str,
default='mnist',
choices=['mnist', 'cifar10'])
parser.add_argument('--data_augmentation',
'-A',
action='store_true',
dest='data_augmentation')
parser.set_defaults(data_augmentation=False)
parser.add_argument('--num_train_examples', type=int, default=None)
parser.add_argument('--label_noise_level', '-n', type=float, default=0.0)
parser.add_argument('--label_noise_type',
type=str,
default='flip',
choices=['flip', 'error', 'cifar10_custom'])
parser.add_argument('--transform_function',
type=str,
default=None,
choices=[None, 'remove_random_chunks'])
parser.add_argument('--clean_validation',
dest='clean_validation',
action='store_true')
parser.set_defaults(clean_validation=False)
parser.add_argument('--remove_prob', type=float, default=0.5)
args = parser.parse_args()
print(args)
# Load data
train_loader, val_loader, test_loader = datasets.load_data_from_arguments(
args)
# Options
optimization_args = {'optimizer': {'name': 'adam', 'lr': 1e-3}}
with open(args.config, 'r') as f:
architecture_args = json.load(f)
revert_normalization = (
lambda x: datasets.revert_normalization(x, train_loader.dataset))
model = VAE(input_shape=train_loader.dataset[0][0].shape,
architecture_args=architecture_args,
device=args.device,
revert_normalization=revert_normalization)
training.train(model=model,
train_loader=train_loader,
val_loader=val_loader,
epochs=args.epochs,
save_iter=args.save_iter,
vis_iter=args.vis_iter,
optimization_args=optimization_args,
log_dir=args.log_dir)