if config.VERBOSE:
state['progress_bar'].update(1)
state['loss_meter'].update(state['loss'].item(), 1)
min_idx = min(state['sample']['batch_anno_idxs'])
batch_indexs = [
idx - min_idx for idx in state['sample']['batch_anno_idxs']
]
sorted_segments = [state['output'][i] for i in batch_indexs]
state['sorted_segments_list'].extend(sorted_segments)
def on_test_end(state):
annotations = state['iterator'].dataset.annotations
state['Rank@N,mIoU@M'], state['miou'] = eval.eval_predictions(
state['sorted_segments_list'], annotations, verbose=False)
if config.VERBOSE:
state['progress_bar'].close()
engine = Engine()
engine.hooks['on_start'] = on_start
engine.hooks['on_forward'] = on_forward
engine.hooks['on_update'] = on_update
engine.hooks['on_end'] = on_end
engine.hooks['on_test_start'] = on_test_start
engine.hooks['on_test_forward'] = on_test_forward
engine.hooks['on_test_end'] = on_test_end
engine.train(network,
iterator('train'),
maxepoch=config.TRAIN.MAX_EPOCH,
optimizer=optimizer,
scheduler=scheduler)
