def from_config(cls, architecture_name, split):
assert split in ['train'] # no augmentation for val and test split
config = load_config(architecture_name + '.yaml')
training_config = load_config('training.yaml')
config.update(training_config)
transformations_config = load_config('random_transformations.yaml')
config.update(transformations_config)
parameters = {**config}
return RandomTransformations(**parameters)
def from_config(cls, architecture_name, phase):
"""
Args:
phase (str): 'training' or 'deployment'
"""
hardnet_config = load_config(architecture_name + '.yaml')
config = load_config(
'training.yaml') if phase == 'training' else load_config(
'deployment.yaml')
hardnet_config.update(config)
hardnet_parameters = {**hardnet_config}
return HarDNet(**hardnet_parameters)
def by_name(cls, architecture_name, phase, verbose=False):
"""Get one of our models by its name.
Args:
architecture_name (str): Currently supported 'hardnet29'
phase (str): 'training' or 'deployment'
verbose (bool): Print some information.
"""
if 'hardnet' in architecture_name:
encoder = HarDNet.from_config(architecture_name, phase)
model_config = load_config(architecture_name + '.yaml')
model_parameters = {**model_config}
model_parameters['encoder'] = encoder
model = Model(**model_parameters)
elif 'densenet' in architecture_name:
encoder = DenseNet.from_config(architecture_name, phase)
model_config = load_config(architecture_name + '.yaml')
model_parameters = {**model_config}
model_parameters['encoder'] = encoder
model = Model(**model_parameters)
else:
raise ValueError("Architechture '{}' is not supported.".format(
architecture_name))
if 'path_to_weights_file' in model_config and model_config[
'path_to_weights_file']:
model = Model.load_weights(
model,
path_to_weights_file=model_config['path_to_weights_file'],
verbose=verbose)
if verbose:
trainable_parameters = filter(
lambda parameter: parameter.requires_grad, model.parameters())
num_trainable_parameters = sum([
np.prod(parameter.size()) for parameter in trainable_parameters
])
print('Number of trainable model parameters: {}'.format(
num_trainable_parameters))
if phase == 'training':
model = model.train()
elif phase == 'deployment':
model = model.eval()
return model
def export_model(architecture_name, path_to_output_file, device, file_type):
assert file_type in ['onnx', 'pt']
print('Exporting model as .{}.'.format(file_type))
try:
model = Model.by_name(architecture_name,
phase='deployment',
verbose=True)
except ValueError:
click.echo(
"Architechture '{}' is not supported.".format(architecture_name),
err=True)
if path_to_output_file is None:
# name similar to architecture_name
path_to_output_file = architecture_name + '.' + file_type
path_to_output_file = Path(path_to_output_file)
if not path_to_output_file.is_absolute():
path_to_output_file = MODELS_DIR / path_to_output_file
click.echo('Writing result to {}'.format(path_to_output_file))
# get input size from configuration file
config = load_config('deployment.yaml')
input_width = config['input_width']
input_height = config['input_height']
input_channels = config['input_channels']
batch_size = config['batch_size']
model = model.to(device)
model.eval()
model.set_deploy(
True) # set flag to add sigmoid, softmax after final layers
dummy_input = torch.randn(batch_size,
input_channels,
input_height,
input_width,
device=device)
model(dummy_input)
if file_type == 'onnx':
torch.onnx.export(
model=model,
args=dummy_input,
f=str(path_to_output_file),
export_params=True,
training=False,
input_names=['input'],
output_names=[
'semantic_output', 'stem_keypoint_output', 'stem_offset_output'
],
verbose=False,
opset_version=7) # maximum opset version supported by tensorrt
elif file_type == 'pt':
traced_script_module = torch.jit.trace(model, dummy_input)
traced_script_module.save(str(path_to_output_file))
def main(args):
cfile = os.path.join(args.target_dir, "config.yaml")
if args.config is not None:
# Already validated by parser
cfile = args.config
else:
try:
utils.validate_file_path(cfile)
except IOError:
cfile = os.path.join(args.target_dir, "config")
utils.validate_file_path(cfile)
config = load_config(cfile)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
transform = REGION_TO_TRANS[config.region]
lon, lat = [transform(i) for i in eg.v1_get_full_grid_lonlat(eg.ML)]
model_path = os.path.join(args.target_dir, "model.pt")
utils.validate_file_path(model_path)
model_class = UNet
if args.legacy_model:
print("Using legacy model")
model_class = UNetLegacy
model = model_class(
in_chan=config.in_chan,
n_classes=config.n_classes,
depth=config.depth,
base_filter_bank_size=config.base_filters,
skip=config.skips,
bndry_dropout=config.bndry_dropout,
bndry_dropout_p=config.bndry_dropout_p,
)
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model = model.to(device)
model_dict = torch.load(model_path)
model_load(model, model_dict)
model.eval()
input_ds = build_input_dataset_form_config(config, False, lat)
batch_size = args.batch_size if args.batch_size > 0 else config.batch_size
input_dl = torch.utils.data.DataLoader(input_ds,
batch_size=batch_size,
shuffle=False,
drop_last=False)
preds, probs = get_predictions(
input_dl,
model,
transform(np.load("../data/masks/ft_esdr_water_mask.npy")),
LABEL_OTHER,
device,
config,
)
# TODO
pred_path = os.path.join(args.target_dir, "pred.npy")
prob_path = os.path.join(args.target_dir, "prob.npy")
print(f"Saving probabilities: '{prob_path}'")
np.save(prob_path, probs)
if args.save_predictions:
print(f"Saving predictions: '{pred_path}'")
np.save(preds, pred_path)
def from_config(cls, architecture_name, split):
assert split in ['train', 'val', 'test']
config = load_config(architecture_name + '.yaml')
# additional parameters for training
training_config = load_config('training.yaml')
config.update(training_config)
# additional parameters for dataset
dataset_config = load_config('sugar_beet_dataset.yaml')
config.update(dataset_config)
dataset_parameters = {**config}
# only use files of the given split as specified in config
# we have 100 in 'val', 100 in 'test' and the rest in 'train'
dataset_parameters['filenames_filter'] = load_config(split +
'_split.yaml')
# hack for evaluation on test set
# if split=='val':
# dataset_parameters['filenames_filter'] = load_config('test_split.yaml')
# else:
# dataset_parameters['filenames_filter'] = load_config(split+'_split.yaml')
# data augmentation for train split
if split == 'train':
random_transformations = RandomTransformations.from_config(
architecture_name, 'train')
else:
random_transformations = None
dataset_parameters['random_transformations'] = random_transformations
return SugarBeetDataset(**dataset_parameters)