def main():
pl.seed_everything(42)
args = parse_arguments()
# Create our dataset.
datamodule = Cifar10DataModule(args.data_root, batch_size=args.batch_size)
datamodule.active_dataset.label_randomly(10)
# Get our heuristic to compute uncertainty.
heuristic = get_heuristic(args.heuristic, shuffle_prop=0.0, reduction='none')
model = VGG16(**vars(args)) # Instantiate VGG16
# Make our PL Trainer
logger = TensorBoardLogger(save_dir=os.path.join('/tmp/', 'logs', 'active'), name='CIFAR10')
trainer = BaalTrainer.from_argparse_args(args,
# The weights of the model will change as it gets
# trained; we need to keep a copy (deepcopy) so that
# we can reset them.
callbacks=[ResetCallback(copy.deepcopy(model.state_dict()))],
dataset=datamodule.active_dataset,
max_epochs=args.training_duration,
logger=logger,
heuristic=heuristic,
ndata_to_label=args.query_size
)
AL_STEPS = 100
for al_step in range(AL_STEPS):
print(f'Step {al_step} Dataset size {len(datamodule.active_dataset)}')
trainer.fit(model, datamodule=datamodule) # Train the model on the labelled set.
trainer.test(model, datamodule=datamodule) # Get test performance.
should_continue = trainer.step(model, datamodule=datamodule) # Label the top-k most uncertain examples.
if not should_continue:
break
def __init__(self, exp_dict):
super().__init__()
self.backbone = models.vgg16(
pretrained=exp_dict["imagenet_pretraining"], progress=True)
num_ftrs = self.backbone.classifier[-1].in_features
self.backbone.classifier[-1] = torch.nn.Linear(num_ftrs,
exp_dict["num_classes"])
self.backbone = patch_module(self.backbone)
self.initial_weights = deepcopy(self.backbone.state_dict())
self.backbone.cuda()
self.batch_size = exp_dict['batch_size']
self.calibrate = exp_dict.get('calibrate', False)
self.learning_epoch = exp_dict['learning_epoch']
self.optimizer = torch.optim.SGD(self.backbone.parameters(),
lr=exp_dict['lr'],
weight_decay=5e-4,
momentum=0.9,
nesterov=True)
self.criterion = CrossEntropyLoss()
shuffle_prop = exp_dict.get('shuffle_prop', 0.0)
max_sample = -1
self.heuristic = get_heuristic(exp_dict['heuristic'],
shuffle_prop=shuffle_prop)
self.wrapper = ModelWrapper(self.backbone, criterion=self.criterion)
self.wrapper.add_metric(
'cls_report',
lambda: ClassificationReport(exp_dict["num_classes"]))
self.wrapper.add_metric('accuracy', lambda: Accuracy())
self.loop = ActiveLearningLoop(None,
self.wrapper.predict_on_dataset,
heuristic=self.heuristic,
ndata_to_label=exp_dict['query_size'],
batch_size=self.batch_size,
iterations=exp_dict['iterations'],
use_cuda=True,
max_sample=max_sample)
self.calib_set = get_dataset('calib', exp_dict['dataset'])
self.valid_set = get_dataset('val', exp_dict['dataset'])
self.calibrator = DirichletCalibrator(
self.wrapper,
exp_dict["num_classes"],
lr=0.001,
reg_factor=exp_dict['reg_factor'],
mu=exp_dict['mu'])
self.active_dataset = None
self.active_dataset_settings = None
def get_data_module(heuristic, data_path):
train_set = datasets.CIFAR10(data_path, train=True, download=True)
test_set = datasets.CIFAR10(data_path, train=False, download=True)
dm = DataModule_.from_datasets(
train_dataset=train_set,
test_dataset=test_set,
train_transform=train_transforms,
test_transform=test_transforms,
# Do not forget to set `predict_transform`,
# this is what we will use for uncertainty estimation!
predict_transform=test_transforms,
batch_size=64,
)
active_dm = ActiveLearningDataModule(
dm,
heuristic=get_heuristic(heuristic),
initial_num_labels=1024,
query_size=100,
val_split=0.0,
)
assert active_dm.has_test, "No test set?"
return active_dm
def test_getting_heuristics(name):
assert isinstance(get_heuristic(name, reduction='mean'), AbstractHeuristic)
def main():
args = parse_args()
batch_size = args.batch_size
use_cuda = torch.cuda.is_available()
hyperparams = vars(args)
pprint(hyperparams)
active_set, test_set = get_datasets(hyperparams['initial_pool'],
hyperparams['data_path'])
# We will use the FocalLoss
criterion = FocalLoss(gamma=2, alpha=0.25)
# Our model is a simple Unet
model = smp.Unet(encoder_name='resnext50_32x4d',
encoder_depth=5,
encoder_weights='imagenet',
decoder_use_batchnorm=False,
classes=len(pascal_voc_ids))
# Add a Dropout layerto use MC-Dropout
add_dropout(model, classes=len(pascal_voc_ids), activation=None)
# This will enable Dropout at test time.
model = MCDropoutModule(model)
# Put everything on GPU.
if use_cuda:
model.cuda()
# Make an optimizer
optimizer = optim.SGD(model.parameters(),
lr=hyperparams["lr"],
momentum=0.9,
weight_decay=5e-4)
# Keep a copy of the original weights
initial_weights = deepcopy(model.state_dict())
# Add metrics
model = ModelWrapper(model, criterion)
model.add_metric('cls_report',
lambda: ClassificationReport(len(pascal_voc_ids)))
# Which heuristic you want to use?
# We will use our custom reduction function.
heuristic = get_heuristic(hyperparams['heuristic'], reduction=mean_regions)
# The ALLoop is in charge of predicting the uncertainty and
loop = ActiveLearningLoop(
active_set,
model.predict_on_dataset_generator,
heuristic=heuristic,
ndata_to_label=hyperparams['n_data_to_label'],
# Instead of predicting on the entire pool, only a subset is used
max_sample=1000,
batch_size=batch_size,
iterations=hyperparams["iterations"],
use_cuda=use_cuda)
acc = []
for epoch in tqdm(range(args.al_step)):
# Following Gal et al. 2016, we reset the weights.
model.load_state_dict(initial_weights)
# Train 50 epochs before sampling.
model.train_on_dataset(active_set, optimizer, batch_size,
hyperparams['learning_epoch'], use_cuda)
# Validation!
model.test_on_dataset(test_set, batch_size, use_cuda)
should_continue = loop.step()
metrics = model.metrics
val_loss = metrics['test_loss'].value
logs = {
"val": val_loss,
"epoch": epoch,
"train": metrics['train_loss'].value,
"labeled_data": active_set._labelled,
"Next Training set size": len(active_set),
'cls_report': metrics['test_cls_report'].value,
}
pprint(logs)
acc.append(logs)
if not should_continue:
break
def main():
args = parse_args()
use_cuda = torch.cuda.is_available()
torch.backends.cudnn.benchmark = True
random.seed(1337)
torch.manual_seed(1337)
if not use_cuda:
print("warning, the experiments would take ages to run on cpu")
hyperparams = vars(args)
active_set, test_set = get_datasets(hyperparams["initial_pool"])
heuristic = get_heuristic(hyperparams["heuristic"],
hyperparams["shuffle_prop"])
criterion = CrossEntropyLoss()
model = vgg16(pretrained=False, num_classes=10)
weights = load_state_dict_from_url(
"https://download.pytorch.org/models/vgg16-397923af.pth")
weights = {k: v for k, v in weights.items() if "classifier.6" not in k}
model.load_state_dict(weights, strict=False)
# change dropout layer to MCDropout
model = patch_module(model)
if use_cuda:
model.cuda()
optimizer = optim.SGD(model.parameters(),
lr=hyperparams["lr"],
momentum=0.9)
# Wraps the model into a usable API.
model = ModelWrapper(model, criterion)
logs = {}
logs["epoch"] = 0
# for prediction we use a smaller batchsize
# since it is slower
active_loop = ActiveLearningLoop(
active_set,
model.predict_on_dataset,
heuristic,
hyperparams.get("query_size", 1),
batch_size=10,
iterations=hyperparams["iterations"],
use_cuda=use_cuda,
)
# We will reset the weights at each active learning step.
init_weights = deepcopy(model.state_dict())
for epoch in tqdm(range(args.epoch)):
# Load the initial weights.
model.load_state_dict(init_weights)
model.train_on_dataset(
active_set,
optimizer,
hyperparams["batch_size"],
hyperparams["learning_epoch"],
use_cuda,
)
# Validation!
model.test_on_dataset(test_set, hyperparams["batch_size"], use_cuda)
metrics = model.metrics
should_continue = active_loop.step()
if not should_continue:
break
val_loss = metrics["test_loss"].value
logs = {
"val": val_loss,
"epoch": epoch,
"train": metrics["train_loss"].value,
"labeled_data": active_set.labelled,
"Next Training set size": len(active_set),
}
print(logs)
def main():
args = parse_args()
use_cuda = torch.cuda.is_available()
torch.backends.cudnn.benchmark = True
random.seed(1337)
torch.manual_seed(1337)
if not use_cuda:
print("warning, the experiments would take ages to run on cpu")
hyperparams = vars(args)
heuristic = get_heuristic(hyperparams['heuristic'],
hyperparams['shuffle_prop'])
model = BertForSequenceClassification.from_pretrained(
pretrained_model_name_or_path=hyperparams["model"])
tokenizer = BertTokenizer.from_pretrained(
pretrained_model_name_or_path=hyperparams["model"])
# In this example we use tokenizer once only in the beginning since it would
# make the whole process faster. However, it is also possible to input tokenizer
# in trainer.
active_set, test_set = get_datasets(hyperparams['initial_pool'], tokenizer)
# change dropout layer to MCDropout
model = patch_module(model)
if use_cuda:
model.cuda()
init_weights = deepcopy(model.state_dict())
training_args = TrainingArguments(
output_dir='/app/baal/results', # output directory
num_train_epochs=hyperparams['learning_epoch'], # total # of training epochs
per_device_train_batch_size=16, # batch size per device during training
per_device_eval_batch_size=64, # batch size for evaluation
weight_decay=0.01, # strength of weight decay
logging_dir='/app/baal/logs', # directory for storing logs
)
# We wrap the huggingface Trainer to create an Active Learning Trainer
model = BaalTransformersTrainer(model=model,
args=training_args,
train_dataset=active_set,
eval_dataset=test_set,
tokenizer=None)
logs = {}
logs['epoch'] = 0
# In this case, nlp data is fast to process and we do NoT need to use a smaller batch_size
active_loop = ActiveLearningLoop(active_set,
model.predict_on_dataset,
heuristic,
hyperparams.get('n_data_to_label', 1),
iterations=hyperparams['iterations'])
for epoch in tqdm(range(args.epoch)):
# we use the default setup of HuggingFace for training (ex: epoch=1).
# The setup is adjustable when BaalHuggingFaceTrainer is defined.
model.train()
# Validation!
eval_metrics = model.evaluate()
# We reorder the unlabelled pool at the frequency of learning_epoch
# This helps with speed while not changing the quality of uncertainty estimation.
should_continue = active_loop.step()
# We reset the model weights to relearn from the new trainset.
model.load_state_dict(init_weights)
model.lr_scheduler = None
if not should_continue:
break
active_logs = {"epoch": epoch,
"labeled_data": active_set._labelled,
"Next Training set size": len(active_set)}
logs = {**eval_metrics, **active_logs}
print(logs)
args.add_argument('--gpus', default=torch.cuda.device_count(), type=int)
args = PIActiveLearningModel.add_model_specific_args(args)
params = args.parse_args()
active_set = ActiveLearningDataset(
CIFAR10(params.data_root,
train=True,
transform=PIModel.train_transform,
download=True),
pool_specifics={'transform': PIModel.test_transform})
active_set.label_randomly(500)
print("Active set length: {}".format(len(active_set)))
print("Pool set length: {}".format(len(active_set.pool)))
heuristic = get_heuristic(params.heuristic)
model = vgg16(pretrained=False, num_classes=10)
weights = load_state_dict_from_url(
'https://download.pytorch.org/models/vgg16-397923af.pth')
weights = {k: v for k, v in weights.items() if 'classifier.6' not in k}
model.load_state_dict(weights, strict=False)
model = PIActiveLearningModel(network=model,
active_dataset=active_set,
hparams=params)
dp = 'dp' if params.gpus > 1 else None
trainer = BaalTrainer(
max_epochs=params.epochs,
default_root_dir=params.data_root,
gpus=params.gpus,
distributed_backend=dp,
