def test_log_dir(self, mock_board): state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))} tboard = TensorBoard(write_epoch_metrics=False) tboard.on_start(state) mock_board.assert_called_once_with(log_dir=os.path.join('./logs', 'Sequential_torchbearer'))
def test_log_dir(self, mock_board): state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))} tboard = TensorBoard(write_epoch_metrics=False) tboard.on_start(state) mock_board.assert_called_once_with( log_dir=os.path.join('./logs', 'Sequential_torchbearer'))
def test_batch_log_dir(self, mock_board): state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0} tboard = TensorBoard(write_batch_metrics=True, write_graph=False, write_epoch_metrics=False) tboard.on_start(state) tboard.on_start_epoch(state) mock_board.assert_called_with(log_dir=os.path.join('./logs', 'Sequential_torchbearer', 'epoch-0'))
def test_batch_writer_closed_on_end_epoch(self, mock_board): mock_board.return_value = Mock() mock_board.return_value.close = Mock() state = {torchbearer.EPOCH: 0} tboard = TensorBoard(write_batch_metrics=True, write_epoch_metrics=False) tboard.on_start_epoch(state) tboard.on_end_epoch({}) mock_board.return_value.close.assert_called_once()
def test_writer_closed_on_end(self, mock_board): mock_board.return_value = Mock() mock_board.return_value.close = Mock() state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))} tboard = TensorBoard(write_epoch_metrics=False) tboard.on_start(state) tboard.on_end({}) mock_board.return_value.close.assert_called_once()
def test_batch_log_dir(self, mock_board, _): state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0} tboard = TensorBoard(write_batch_metrics=True, write_graph=False, write_epoch_metrics=False) tboard.on_start(state) tboard.on_start_epoch(state) tboard.on_end_epoch(state) tboard.on_end(state) mock_board.assert_called_with(log_dir=os.path.join('./logs', 'Sequential_torchbearer', 'epoch-0'))
def test_epoch_metrics(self, mock_board): mock_board.return_value = Mock() mock_board.return_value.add_scalar = Mock() state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0, torchbearer.METRICS: {'test': 1}} tboard = TensorBoard(write_batch_metrics=False, write_epoch_metrics=True) tboard.on_start(state) tboard.on_end_epoch(state) mock_board.return_value.add_scalar.assert_called_once_with('epoch/test', 1, 0)
def test_batch_log_dir_visdom(self, mock_visdom, mock_writer, _): state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0, torchbearer.METRICS: {'test': 1}, torchbearer.BATCH: 0} tboard = TensorBoard(visdom=True, write_batch_metrics=True, write_graph=False, write_epoch_metrics=False) tboard.on_start(state) tboard.on_start_epoch(state) tboard.on_end_epoch(state) tboard.on_end(state) self.assertTrue(mock_visdom.call_args[1]['log_to_filename'] == os.path.join('./logs', 'Sequential_torchbearer', 'epoch', 'log.log'))
def test_batch_writer_closed_on_end_epoch_visdom(self, mock_visdom, mock_writer, _): mock_writer.return_value = Mock() mock_writer.return_value.close = Mock() state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0} tboard = TensorBoard(visdom=True, write_batch_metrics=True, write_epoch_metrics=False) tboard.on_start(state) tboard.on_start_epoch(state) tboard.on_end_epoch({}) tboard.on_end(state) self.assertTrue(mock_writer.return_value.close.call_count == 2)
def test_epoch_metrics_visdom(self, mock_visdom, mock_writer, _): mock_writer.return_value = Mock() mock_writer.return_value.add_scalar = Mock() state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0, torchbearer.METRICS: {'test': 1}} tboard = TensorBoard(visdom=True, write_batch_metrics=False, write_epoch_metrics=True) tboard.on_start(state) tboard.on_start_epoch(state) tboard.on_end_epoch(state) mock_writer.return_value.add_scalar.assert_called_once_with('test', 1, 0, main_tag='epoch') tboard.on_end(state)
def test_write_graph(self, mock_rand, mock_board): mock_board.return_value = Mock() mock_board.return_value.add_graph = Mock() mock_rand.return_value = 1 state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.X: torch.zeros(1, 1, 9, 9)} tboard = TensorBoard(write_epoch_metrics=False) tboard.on_start(state) tboard.on_sample(state) mock_rand.assert_called_once_with(state[torchbearer.X].size(), requires_grad=False) mock_board.return_value.add_graph.assert_called_once() self.assertEqual(str(state[torchbearer.MODEL]), str(mock_board.return_value.add_graph.call_args_list[0][0][0])) self.assertNotEqual(state[torchbearer.MODEL], mock_board.return_value.add_graph.call_args_list[0][0][0])
def test_batch_writer_closed_on_end_epoch(self, mock_board): mock_board.return_value = Mock() mock_board.return_value.close = Mock() state = { torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.EPOCH: 0 } tboard = TensorBoard(write_batch_metrics=True, write_epoch_metrics=False) tboard.on_start(state) tboard.on_start_epoch(state) tboard.on_end_epoch({}) mock_board.return_value.close.assert_called_once() tboard.on_end(state)
def test_visdom_error(self, mock_visdom, mock_writer, _, mock_warn): import unittest.mock orig_import = __import__ def import_mock(name, *args): if name == 'visdom': raise ImportError('Test error') return orig_import(name, *args) with unittest.mock.patch('builtins.__import__', side_effect=import_mock): state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))} tboard = TensorBoard(visdom=True, write_epoch_metrics=False) tboard.on_start(state) self.assertEqual(mock_warn.call_count, 1)
def test_state_dict(self): self.callback_1.state_dict = Mock(return_value='test_1') self.callback_2.state_dict = Mock(return_value='test_2') state = self.list.state_dict() self.assertEqual(self.callback_1.state_dict.call_count, 1) self.assertEqual(self.callback_2.state_dict.call_count, 1) self.assertEqual(state[CallbackList.CALLBACK_STATES][0], 'test_1') self.assertEqual(state[CallbackList.CALLBACK_STATES][1], 'test_2') self.assertEqual(state[CallbackList.CALLBACK_TYPES][0], Tqdm().__class__) self.assertEqual(state[CallbackList.CALLBACK_TYPES][1], TensorBoard().__class__)
def test_write_graph(self, mock_rand, mock_board, _): mock_board.return_value = Mock() mock_board.return_value.add_graph = Mock() mock_rand.return_value = 1 state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3)), torchbearer.X: torch.zeros(1, 1, 9, 9)} tboard = TensorBoard(write_epoch_metrics=False) tboard.on_start(state) tboard.on_sample(state) tboard.on_end(state) mock_rand.assert_called_once_with(state[torchbearer.X].size(), requires_grad=False) self.assertEqual(mock_board.return_value.add_graph.call_count, 1) self.assertEqual(str(state[torchbearer.MODEL]), str(mock_board.return_value.add_graph.call_args_list[0][0][0])) self.assertNotEqual(state[torchbearer.MODEL], mock_board.return_value.add_graph.call_args_list[0][0][0])
def test_batch_metrics(self, mock_board): mock_board.return_value = Mock() mock_board.return_value.add_scalar = Mock() state = {torchbearer.EPOCH: 0, torchbearer.METRICS: {'test': 1}, torchbearer.BATCH: 0} tboard = TensorBoard(write_batch_metrics=True, write_epoch_metrics=False) tboard.on_start_epoch(state) tboard.on_step_training(state) mock_board.return_value.add_scalar.assert_called_once_with('batch/test', 1, 0) mock_board.return_value.add_scalar.reset_mock() tboard.on_step_validation(state) mock_board.return_value.add_scalar.assert_called_once_with('batch/test', 1, 0)
def test_writer_closed_on_end(self, mock_board, _): mock_board.return_value = Mock() mock_board.return_value.close = Mock() state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))} tboard = TensorBoard(write_epoch_metrics=False) tboard.on_start(state) tboard.on_end({}) self.assertEqual(mock_board.return_value.close.call_count, 1)
def test_writer_closed_on_end_visdom_visdom(self, mock_visdom, mock_writer, _): mock_writer.return_value = Mock() mock_writer.return_value.close = Mock() state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))} tboard = TensorBoard(visdom=True) tboard.on_start(state) tboard.on_end({}) self.assertEqual(mock_writer.return_value.close.call_count, 1)
def test_writer_closed_on_end_visdom(self, mock_visdom, mock_writer, _): mock_writer.return_value = Mock() mock_writer.return_value.close = Mock() state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))} tboard = TensorBoard(visdom=True, write_epoch_metrics=False) tboard.on_start(state) tboard.on_end({}) mock_writer.return_value.close.assert_called_once()
def test_log_dir_visdom(self, mock_visdom, mock_writer, _): state = {torchbearer.MODEL: nn.Sequential(nn.Conv2d(3, 3, 3))} mock_writer.__delete__ = Mock() tboard = TensorBoard(visdom=True, write_epoch_metrics=False) tboard.on_start(state) tboard.on_end(state) self.assertEqual(mock_visdom.call_count, 1) self.assertTrue(mock_visdom.call_args[1]['log_to_filename'] == os.path.join('./logs', 'Sequential_torchbearer', 'log.log'))
def test_batch_metrics(self, mock_board): mock_board.return_value = Mock() mock_board.return_value.add_scalar = Mock() state = { torchbearer.EPOCH: 0, torchbearer.METRICS: { 'test': 1 }, torchbearer.BATCH: 0 } tboard = TensorBoard(write_batch_metrics=True, write_epoch_metrics=False) tboard.on_start_epoch(state) tboard.on_step_training(state) mock_board.return_value.add_scalar.assert_called_once_with( 'batch/test', 1, 0) mock_board.return_value.add_scalar.reset_mock() tboard.on_step_validation(state) mock_board.return_value.add_scalar.assert_called_once_with( 'batch/test', 1, 0)
elif state[tb.MODEL].x < -1: state[tb.MODEL].x.data.fill_(-1) training_steps = 6000000 model = Online() optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99]) tbtrial = tb.Trial(model, optim, loss, [est()], pass_state=True, callbacks=[ greedy_update, TensorBoard(comment='adam', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False) ]) tbtrial.for_train_steps(training_steps).run() model = Online() optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99], amsgrad=True) tbtrial = tb.Trial(model, optim, loss, [est()], pass_state=True, callbacks=[ greedy_update,
nn.ReLU(), nn.Conv2d(16, 32, stride=2, kernel_size=3), nn.BatchNorm2d(32), nn.ReLU(), nn.Conv2d(32, 64, stride=2, kernel_size=3), nn.BatchNorm2d(64), nn.ReLU() ) self.classifier = nn.Linear(576, 10) def forward(self, x): x = self.convs(x) x = x.view(-1, 576) return self.classifier(x) model = SimpleModel() optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001) loss = nn.CrossEntropyLoss() import torchbearer from torchbearer import Trial from torchbearer.callbacks import TensorBoard torchbearer_trial = Trial(model, optimizer, loss, metrics=['acc', 'loss'], callbacks=[TensorBoard(write_batch_metrics=True)]).to('cuda') torchbearer_trial.with_generators(train_generator=traingen, val_generator=valgen, test_generator=testgen) torchbearer_trial.run(epochs=10) torchbearer_trial.evaluate(data_key=torchbearer.TEST_DATA)
# test_set = CIFAR10('./data', train=False, download=True, transform=transforms.Compose([transform_test])) test_gen = torch.utils.data.DataLoader(test_set, pin_memory=True, batch_size=128, shuffle=False, num_workers=4) optimizer = optim.SGD(net.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4) from torchbearer.callbacks import MultiStepLR, TensorBoard from torchbearer import Trial from datetime import datetime current_time = datetime.now().strftime('%b%d_%H-%M-%S') trial = Trial(net, optimizer, nn.CrossEntropyLoss(), metrics=['acc', 'loss'], callbacks=[ UnpackState(), TensorBoard(write_graph=False, comment=current_time), MultiStepLR([100, 150]) ]) trial.with_generators(train_generator=train_gen, val_generator=test_gen).to('cuda') trial.run(200, verbose=1)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) if (trainset is not None) and (args.dataset not in nlp_data) else trainset valloader = torch.utils.data.DataLoader(valset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) if (valset is not None) and (args.dataset not in nlp_data) else valset testloader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) if (args.dataset not in nlp_data) else testset print('==> Building model..') net = get_model(args, classes, nc) net = nn.DataParallel(net) if args.parallel else net optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) if (args.dataset in nlp_data) or ('modelnet' in args.dataset): optimizer = optim.Adam(net.parameters(), lr=args.lr) print('==> Setting up callbacks..') current_time = datetime.now().strftime('%b%d_%H-%M-%S') + "-run-" + str(args.run_id) tboard = TensorBoard(write_graph=False, comment=current_time, log_dir=args.log_dir) tboardtext = TensorBoardText(write_epoch_metrics=False, comment=current_time, log_dir=args.log_dir) @torchbearer.callbacks.on_start def write_params(_): params = vars(args) params['schedule'] = str(params['schedule']) df = pd.DataFrame(params, index=[0]).transpose() tboardtext.get_writer(tboardtext.log_dir).add_text('params', df.to_html(), 1) modes = { 'fmix': FMix(decay_power=args.f_decay, alpha=args.alpha, size=size, max_soft=0, reformulate=args.reformulate), 'mixup': RMixup(args.alpha, reformulate=args.reformulate), 'cutmix': CutMix(args.alpha, classes, True),
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001) loss = nn.CrossEntropyLoss() from torchbearer import Trial from torchbearer.callbacks import TensorBoard torchbearer_trial = Trial(model, optimizer, loss, metrics=['acc', 'loss'], callbacks=[ TensorBoard(visdom=True, write_graph=True, write_batch_metrics=True, batch_step_size=10, write_epoch_metrics=True) ]).to('cuda') torchbearer_trial.with_generators(train_generator=traingen, val_generator=valgen) torchbearer_trial.run(epochs=5) import torchbearer.callbacks.tensor_board as tensorboard tensorboard.VisdomParams.ENV = 'Test' torchbearer_trial = Trial(model, optimizer, loss, metrics=['acc', 'loss'], callbacks=[
return state[tb.MODEL].x.data @tb.callbacks.on_step_training def greedy_update(state): if state[tb.MODEL].x > 1: state[tb.MODEL].x.data.fill_(1) elif state[tb.MODEL].x < -1: state[tb.MODEL].x.data.fill_(-1) training_steps = 6000000 model = Online() optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99]) tbtrial = tb.Trial(model, optim, loss, [est()], pass_state=True, callbacks=[greedy_update, TensorBoard(comment='adam', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)]) tbtrial.for_train_steps(training_steps).run() model = Online() optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99], amsgrad=True) tbtrial = tb.Trial(model, optim, loss, [est()], pass_state=True, callbacks=[greedy_update, TensorBoard(comment='amsgrad', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)]) tbtrial.for_train_steps(training_steps).run() model = Stochastic() optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99]) tbtrial = tb.Trial(model, optim, loss, [est()], callbacks=[greedy_update, TensorBoard(comment='adam', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)]) tbtrial.for_train_steps(training_steps).run() model = Stochastic() optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99], amsgrad=True) tbtrial = tb.Trial(model, optim, loss, [est()], callbacks=[greedy_update, TensorBoard(comment='amsgrad', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)])