def preprocess_dataset(data_name):
input_dir = get_data_dir(cfg, data_name)
print('\nDataset:', input_dir)
if not os.path.isdir(input_dir):
print('Not found')
return
# Create the output directory
output_name = data_name + '.' + WORKER_UID
output_dir = os.path.join(cfg.preproc_dir, output_name)
os.makedirs(output_dir)
# Preprocess image sample groups
sample_groups = get_image_sample_groups(input_dir, cfg.features)
tza_filename = os.path.join(output_dir, 'images.tza')
samples = []
with tza.Writer(tza_filename) as output_tza:
for _, input_names, target_name in sample_groups:
if target_name:
samples += preprocess_sample_group(input_dir, output_tza,
input_names,
target_name)
# Save the samples in the dataset
samples_filename = os.path.join(output_dir, 'samples.json')
save_json(samples_filename, samples)
# Save the config
save_config(output_dir, cfg)
def main():
# Parse the command line arguments
cfg = parse_args(
description=
'Exports a training result to the runtime model weights format (TZA).')
# Initialize the PyTorch device
device = init_device(cfg)
# Load the checkpoint
checkpoint = load_checkpoint(cfg, device, cfg.checkpoint)
epoch = checkpoint['epoch']
model_state = checkpoint['model_state']
# Save the weights to a TZA file
output_filename = os.path.join(get_result_dir(cfg), cfg.result)
if cfg.checkpoint:
output_filename += '_%d' % epoch
output_filename += '.tza'
print('Saving weights:', output_filename)
with tza.Writer(output_filename) as output_file:
for name, value in model_state.items():
tensor = value.cpu().numpy()
print(name, tensor.shape)
if name.endswith('.weight') and len(value.shape) == 4:
layout = 'oihw'
elif len(value.shape) == 1:
layout = 'x'
else:
error('unknown state value')
output_file.write(name, tensor, layout)
def export_weights(cfg):
# Initialize the PyTorch device
device = init_device(cfg)
# Load the checkpoint
result_dir = get_result_dir(cfg)
checkpoint = load_checkpoint(result_dir, device, cfg.num_epochs)
epoch = checkpoint['epoch']
model_state = checkpoint['model_state']
print('Epoch:', epoch)
# Save the weights to a TZA file
if cfg.output:
output_filename = cfg.output
else:
output_filename = os.path.join(result_dir, cfg.result)
if cfg.num_epochs:
output_filename += '_%d' % epoch
output_filename += '.tza'
print('Output:', output_filename)
print()
with tza.Writer(output_filename) as output_file:
for name, value in model_state.items():
tensor = value.cpu().numpy()
print(name, tensor.shape)
if name.endswith('.weight') and len(value.shape) == 4:
layout = 'oihw'
elif len(value.shape) == 1:
layout = 'x'
else:
error('unknown state value')
output_file.write(name, tensor, layout)
def preprocess_dataset(data_name):
input_dir = get_data_dir(cfg, data_name)
output_dir = get_preproc_data_dir(cfg, data_name)
print('Preprocessing dataset:', input_dir)
# Check whether the dataset exists
if not os.path.isdir(input_dir):
print('Does not exist')
return
# Create the output directory if it doesn't exist
if os.path.isdir(output_dir):
if cfg.clean:
shutil.rmtree(output_dir)
else:
print('Skipping, already preprocessed')
return
os.makedirs(output_dir)
# Save the config
save_config(output_dir, cfg)
# Preprocess image sample groups
sample_groups = get_image_sample_groups(input_dir, cfg.features)
tza_filename = os.path.join(output_dir, 'images.tza')
samples = []
with tza.Writer(tza_filename) as output_tza:
for _, input_names, target_name in sample_groups:
if target_name:
samples += preprocess_sample_group(input_dir, output_tza,
input_names,
target_name)
# Save the samples in the dataset
samples_filename = os.path.join(output_dir, 'samples.json')
save_json(samples_filename, samples)