def get_test_image_preprocessor(batch_size, params):
"""Returns the preprocessing.TestImagePreprocessor that should be injected.
Returns None if no preprocessor should be injected.
Args:
batch_size: The batch size across all GPUs.
params: BenchmarkCNN's parameters.
Returns:
Returns the preprocessing.TestImagePreprocessor that should be injected.
Raises:
ValueError: Flag --fake_input is an invalid value.
"""
if FLAGS.fake_input == 'none':
return None
elif FLAGS.fake_input == 'zeros_and_ones':
half_batch_size = batch_size // 2
images = np.zeros((batch_size, 227, 227, 3), dtype=np.float32)
images[half_batch_size:, :, :, :] = 1
labels = np.array([0] * half_batch_size + [1] * half_batch_size,
dtype=np.int32)
preprocessor = preprocessing.TestImagePreprocessor(
227, 227, batch_size, params.num_gpus,
benchmark_cnn.get_data_type(params))
preprocessor.set_fake_data(images, labels)
preprocessor.expected_subset = 'validation' if params.eval else 'train'
return preprocessor
else:
raise ValueError('Invalid --fake_input: %s' % FLAGS.fake_input)
def _run_benchmark_cnn_with_fake_images(self, params, images, labels):
logs = []
benchmark_cnn.log_fn = _print_and_add_to_list(logs)
bench = benchmark_cnn.BenchmarkCNN(params)
bench.image_preprocessor = preprocessing.TestImagePreprocessor(
227, 227, params.batch_size * params.num_gpus, params.num_gpus,
benchmark_cnn.get_data_type(params))
bench.dataset._queue_runner_required = True
bench.image_preprocessor.set_fake_data(images, labels)
bench.image_preprocessor.expected_subset = ('validation' if params.eval
else 'train')
bench.run()
return logs
def train(config):
# SETUP
components = setup.setup(config)
image_info = components["image_info"]
heads_info = components["heads_info"]
output_files = components["output_files"]
state_folder = components["state_folder"]
image_folder = components["image_folder"]
label_folder = components["label_folder"]
net = components["net"]
# FORCE RESTART
if config.output.force_training_restart:
output_files.clear_output()
# OPTIMIZER
optimizer = torch.optim.Adam(net.parameters(), lr=config.optimizer.learning_rate)
# MODEL
state_folder = output_files.get_sub_root(output_files.STATE)
if Model.exists(state_folder):
# LOAD EXISTING
model = Model.load(state_folder, net=net, optimizer=optimizer)
else:
# CREATE NEW
# LOSS
iid_loss = loss.IIDLoss(
heads=heads_info.order,
output_files=output_files,
do_render=config.output.rendering.enabled,
**config.training.loss
)
# STATISTICS
epoch_stats = utils.EpochStatistics(
limit=config.training.num_epochs, output_files=output_files
)
model = Model(
state_folder=state_folder,
heads_info=heads_info,
net=net,
optimizer=optimizer,
loss_fn=iid_loss,
epoch_statistics=epoch_stats,
)
# PREPROCESSING
# transformation
transformation = pre.Transformation(**config.transformations)
# label mapping
LABEL_FILTERS = {"CocoFewLabels": cocostuff.CocoFewLabels}
if (
"label_filter" in config.dataset
and config.dataset.label_filter.name in LABEL_FILTERS
):
label_filter = LABEL_FILTERS[config.dataset.label_filter.name](
class_count=heads_info.class_count, **config.dataset.label_filter.parameters
)
label_mapper = pre.LabelMapper(mapping_function=label_filter.apply)
else:
print("unable to find label mapper, using identity mapping")
label_mapper = pre.LabelMapper()
# general preprocessing
preprocessing = pre.TransformPreprocessing(
transformation=transformation,
image_info=image_info,
label_mapper=label_mapper,
**config.preprocessor
)
# TRAIN DATALOADER
train_prep = pre.TrainImagePreprocessor(
image_info=image_info,
preprocessing=preprocessing,
output_files=output_files,
do_render=config.output.rendering.enabled,
render_limit=config.output.rendering.limit,
)
train_dataset = data.ImageFolderDataset(
image_folder=image_folder,
preprocessor=train_prep,
extensions=config.dataset.extensions,
label_folder=label_folder,
)
train_dataloader = data.TrainDataLoader(
dataset=train_dataset,
batch_size=config.training.batch_size,
shuffle=config.training.shuffle,
)
# TEST DATALOADER
test_prep = pre.TestImagePreprocessor(
image_info=image_info,
preprocessing=preprocessing,
output_files=output_files,
do_render=config.output.rendering.enabled,
render_limit=config.output.rendering.limit,
)
test_dataset = data.ImageFolderDataset(
image_folder=image_folder,
preprocessor=test_prep,
extensions=config.dataset.extensions,
label_folder=label_folder,
)
test_dataloader = data.TestDataLoader(
dataset=test_dataset, batch_size=config.training.batch_size
)
# DATALOADERS
# TODO link this to arch heads
dataloaders = {
"A": train_dataloader,
"B": train_dataloader,
"map_assign": test_dataloader,
"map_test": test_dataloader,
}
model.train(loaders=dataloaders)