def test_end_to_end(self):
with torchbearer.no_grad():
model = torch.nn.Linear(10, 5)
callback = imaging.ClassAppearanceModel(nclasses=5, input_size=(10), steps=1)
state = {torchbearer.EPOCH: 0, torchbearer.MODEL: model, torchbearer.DEVICE: 'cpu',
torchbearer.DATA_TYPE: torch.float32}
callback.on_batch(state)
def test_random_targets(self, randint):
randint.return_value = torch.ones(1, 1) * 6
callback = imaging.ClassAppearanceModel(nclasses=10, input_size=(1, 1, 1))
state = {torchbearer.DEVICE: 'cpu'}
targets_hot = callback._targets_hot(state)
self.assertTrue(targets_hot.sum() == 1)
self.assertTrue(targets_hot[0][6] == 1)
def test_targets_to_key(self):
callback = imaging.ClassAppearanceModel(nclasses=10, input_size=(1, 1, 1), target=5)
callback = callback.target_to_key('test')
state = {torchbearer.DEVICE: 'cpu'}
targets_hot = callback._targets_hot(state)
self.assertTrue(targets_hot.sum() == 1)
self.assertTrue(targets_hot[0][5] == 1)
self.assertTrue(state['test'] == 5)
def test_on_batch(self, _, wrapper, loss):
wrapper().input_image = torch.nn.Parameter(torch.zeros(10))
factory = MagicMock()
callback = imaging.ClassAppearanceModel(nclasses=10, input_size=(1, 1, 1), optimizer_factory=factory)
model = MagicMock()
callback.on_batch({torchbearer.EPOCH: 0, torchbearer.MODEL: model, torchbearer.DEVICE: 'cpu', torchbearer.DATA_TYPE: torch.float32})
loss.assert_called_once_with(torchbearer.PREDICTION, ANY, 0.01)
self.assertTrue(model.eval.call_count == 1)
self.assertTrue(model.train.call_count == 1)
self.assertTrue(factory.call_count == 1)
self.assertTrue(next(iter(factory.call_args[0][0])) is wrapper().input_image)
def test_targets_hot(self):
callback = imaging.ClassAppearanceModel(nclasses=10, input_size=(1, 1, 1), target=5)
targets_hot = callback._targets_hot({torchbearer.DEVICE: 'cpu'})
self.assertTrue(targets_hot.sum() == 1)
self.assertTrue(targets_hot[0][5] == 1)
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.net = torchvision.models.googlenet(True)
def forward(self, input):
if input is not None:
return self.net(input)
model = Model()
from torchbearer import Trial
import torchbearer.callbacks.imaging as imaging
trial = Trial(model,
callbacks=[
imaging.ClassAppearanceModel(
1000, (3, 224, 224),
steps=10000,
target=951,
transform=inv_normalize).on_val().to_file('lemon.png'),
imaging.ClassAppearanceModel(
1000, (3, 224, 224),
steps=10000,
target=968,
transform=inv_normalize).on_val().to_file('cup.png')
])
trial.for_val_steps(1).to('cuda')
trial.evaluate()