def setUp(self):
self.model = DummyModel()
self.criterion = nn.CrossEntropyLoss()
self.wrapper = ModelWrapper(self.model, self.criterion)
self.optim = torch.optim.SGD(self.wrapper.get_params(), 0.01)
self.dataset = DummyDataset()
self.calibrator = DirichletCalibrator(self.wrapper, 2, lr=0.001, reg_factor=0.001)
def test_calibration_integration():
transform_pipeline = Compose([Resize((64, 64)), ToTensor()])
cifar10_train = DummyDataset(transform_pipeline)
cifar10_test = DummyDataset(transform_pipeline)
# we don't create different trainset for calibration since the goal is not
# to calibrate
al_dataset = ActiveLearningDataset(
cifar10_train, pool_specifics={'transform': transform_pipeline})
al_dataset.label_randomly(10)
use_cuda = False
model = vgg.vgg16(pretrained=False, num_classes=10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=0.001,
momentum=0.9,
weight_decay=0.0005)
wrapper = ModelWrapper(model, criterion)
calibrator = DirichletCalibrator(wrapper=wrapper,
num_classes=10,
lr=0.001,
reg_factor=0.01)
for step in range(2):
wrapper.train_on_dataset(al_dataset,
optimizer=optimizer,
batch_size=10,
epoch=1,
use_cuda=use_cuda,
workers=0)
wrapper.test_on_dataset(cifar10_test,
batch_size=10,
use_cuda=use_cuda,
workers=0)
before_calib_param = list(
map(lambda x: x.clone(), wrapper.model.parameters()))
calibrator.calibrate(al_dataset,
cifar10_test,
batch_size=10,
epoch=5,
use_cuda=use_cuda,
double_fit=False,
workers=0)
after_calib_param = list(map(lambda x: x.clone(), model.parameters()))
assert all([
np.allclose(i.detach(), j.detach())
for i, j in zip(before_calib_param, after_calib_param)
])
assert len(list(wrapper.model.modules())) < len(
list(calibrator.calibrated_model.modules()))
def setUp(self):
# self.model = nn.Sequential(
# nn.Linear(10, 8), nn.ReLU(), nn.Dropout(), nn.Linear(8, 1), nn.Sigmoid()
# )
self.model = DummyModel()
self.criterion = nn.BCEWithLogitsLoss()
self.wrapper = ModelWrapper(self.model, self.criterion)
self.optim = torch.optim.SGD(self.wrapper.get_params(), 0.01)
self.dataset = DummyDataset()
def test_integration():
transform_pipeline = Compose([Resize((64, 64)), ToTensor()])
cifar10_train = DummyDataset(transform_pipeline)
cifar10_test = DummyDataset(transform_pipeline)
al_dataset = ActiveLearningDataset(
cifar10_train, pool_specifics={'transform': transform_pipeline})
al_dataset.label_randomly(10)
use_cuda = False
model = vgg.vgg16(pretrained=False, num_classes=10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=0.001,
momentum=0.9,
weight_decay=0.0005)
# We can now use BaaL to create the active learning loop.
model = ModelWrapper(model, criterion)
# We create an ActiveLearningLoop that will automatically label the most uncertain samples.
# In this case, we use the widely used BALD heuristic.
active_loop = ActiveLearningLoop(al_dataset,
model.predict_on_dataset,
heuristic=heuristics.BALD(),
ndata_to_label=10,
batch_size=10,
iterations=10,
use_cuda=use_cuda,
workers=4)
# We're all set!
num_steps = 10
for step in range(num_steps):
old_param = list(map(lambda x: x.clone(), model.model.parameters()))
model.train_on_dataset(al_dataset,
optimizer=optimizer,
batch_size=10,
epoch=5,
use_cuda=use_cuda,
workers=2)
model.test_on_dataset(cifar10_test,
batch_size=10,
use_cuda=use_cuda,
workers=2)
if not active_loop.step():
break
new_param = list(map(lambda x: x.clone(), model.model.parameters()))
assert any([
not np.allclose(i.detach(), j.detach())
for i, j in zip(old_param, new_param)
])
assert step == 4 # 10 + (4 * 10) = 50, so it stops at iterations 4
def pred_without_dropout(replicate_in_memory):
self.wrapper = ModelWrapper(
self.model,
self.criterion,
replicate_in_memory=replicate_in_memory)
# Dropout is not active in eval
self.wrapper.eval()
preds = torch.stack([
self.wrapper.predict_on_batch(input, iterations=1, cuda=False)
for _ in range(10)
]).view(10, -1)
assert torch.allclose(torch.mean(preds, 0), preds[0])
def pred_with_dropout(replicate_in_memory):
self.wrapper = ModelWrapper(
self.model,
self.criterion,
replicate_in_memory=replicate_in_memory)
self.wrapper.train()
# Dropout make the pred changes
preds = torch.stack([
self.wrapper.predict_on_batch(input, iterations=1, cuda=False)
for _ in range(10)
]).view(10, -1)
assert not torch.allclose(torch.mean(preds, 0), preds[0])
def setUp(self):
class MultiOutModel(nn.Module):
def __init__(self):
super().__init__()
self.model = DummyModel()
def forward(self, x):
return [self.model(x)] * 2
self._crit = nn.MSELoss()
self.criterion = lambda x, y: self._crit(x[0], y) + self._crit(x[1], y)
self.model = MultiOutModel()
self.wrapper = ModelWrapper(self.model, self.criterion)
self.optim = torch.optim.SGD(self.wrapper.get_params(), 0.01)
self.dataset = DummyDataset()
def test_multi_input_model():
class MultiInModel(nn.Module):
def __init__(self):
super().__init__()
self.model = DummyModel()
def forward(self, x):
# We get two inputs
x1, x2 = x
# We merge those inputs
return self.model(x1) + self.model(x2)
model = MultiInModel()
wrapper = ModelWrapper(model, None)
dataset = DummyDataset(n_in=2)
assert len(dataset[0]) == 2
b = next(iter(DataLoader(dataset, 15, False)))[0]
l = wrapper.predict_on_batch(b, iterations=10, cuda=False)
assert l.shape[0] == 15 and l.shape[-1] == 10
def test_predict_on_batch(self):
self.wrapper.eval()
input = torch.stack((self.dataset[0][0], self.dataset[1][0]))
# iteration == 1
pred = self.wrapper.predict_on_batch(input, 1, False)
assert pred[0].size() == (2, 1, 1)
# iterations > 1
pred = self.wrapper.predict_on_batch(input, 10, False)
assert pred[0].size() == (2, 1, 10)
# iteration == 1
self.wrapper = ModelWrapper(self.model,
self.criterion,
replicate_in_memory=False)
pred = self.wrapper.predict_on_batch(input, 1, False)
assert pred[0].size() == (2, 1, 1)
# iterations > 1
pred = self.wrapper.predict_on_batch(input, 10, False)
assert pred[0].size() == (2, 1, 10)
def test_predict_on_batch(self):
self.wrapper.eval()
input = torch.randn([2, 3, 10, 10])
# iteration == 1
pred = self.wrapper.predict_on_batch(input, 1, False)
assert pred.size() == (2, 1, 1)
# iterations > 1
pred = self.wrapper.predict_on_batch(input, 10, False)
assert pred.size() == (2, 1, 10)
# iteration == 1
self.wrapper = ModelWrapper(self.model,
self.criterion,
replicate_in_memory=False)
pred = self.wrapper.predict_on_batch(input, 1, False)
assert pred.size() == (2, 1, 1)
# iterations > 1
pred = self.wrapper.predict_on_batch(input, 10, False)
assert pred.size() == (2, 1, 10)
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('n_data_to_label', 1),
batch_size=10,
iterations=hyperparams['iterations'],
use_cuda=use_cuda)
for epoch in tqdm(range(args.epoch)):
model.train_on_dataset(active_set, optimizer,
hyperparams["batch_size"], 1, use_cuda)
# Validation!
model.test_on_dataset(test_set, hyperparams["batch_size"], use_cuda)
metrics = model.metrics
if epoch % hyperparams['learning_epoch'] == 0:
should_continue = active_loop.step()
model.reset_fcs()
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)