def _init_graph(config, with_dataset=False):
set_seed(config.get('seed', int.from_bytes(os.urandom(4),
byteorder='big')))
n_gpus = get_num_gpus()
logging.info('Number of GPUs detected: {}'.format(n_gpus))
model = get_model(config['model']['name'])(data={},
n_gpus=n_gpus,
**config['model'])
model.__enter__()
if with_dataset:
yield model
else:
yield model
model.__exit__()
tf.reset_default_graph()
def _init_graph(config, with_dataset=False):
set_seed(config.get('seed', int.from_bytes(os.urandom(4),
byteorder='big')))
n_gpus = len(os.environ['CUDA_VISIBLE_DEVICES'].split(','))
logging.info('Number of GPUs detected: {}'.format(n_gpus))
dataset = get_dataset(config['data']['name'])(**config['data'])
model = get_model(config['model']['name'])(
data={} if with_dataset else dataset.get_tf_datasets(),
n_gpus=n_gpus,
**config['model'])
model.__enter__()
if with_dataset:
yield model, dataset
else:
yield model
model.__exit__()
tf.reset_default_graph()
from superpoint.models import get_model # noqa: E402
from superpoint.settings import EXPER_PATH # noqa: E402
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('config', type=str)
parser.add_argument('export_name', type=str)
args = parser.parse_args()
export_name = args.export_name
with open(args.config, 'r') as f:
config = yaml.load(f)
config['model']['data_format'] = 'channels_last'
export_root_dir = Path(EXPER_PATH, 'saved_models')
export_root_dir.mkdir(parents=True, exist_ok=True)
export_dir = Path(export_root_dir, export_name)
checkpoint_path = Path(EXPER_PATH, export_name)
with get_model(config['model']['name'])(data_shape={
'image': [None, None, None, 1]
},
**config['model']) as net:
net.load(str(checkpoint_path))
tf.saved_model.simple_save(net.sess,
str(export_dir),
inputs=net.pred_in,
outputs=net.pred_out)