dmlc--dgl
Konstantinos Vandikas
eca9b8ef08
* allow for configuring default_dir * allow for using DGLDEFAULTDIR environment variable * Update env_var.rst Co-authored-by: Quan (Andy) Gan <coin2028@hotmail.com> Co-authored-by: Jinjing Zhou <VoVAllen@users.noreply.github.com>
23 行
1021 B
Python
23 行
1021 B
Python
import argparse
|
|
import os
|
|
import json
|
|
|
|
def set_default_backend(default_dir, backend_name):
|
|
if not os.path.exists(default_dir):
|
|
os.makedirs(default_dir)
|
|
config_path = os.path.join(default_dir, 'config.json')
|
|
with open(config_path, "w") as config_file:
|
|
json.dump({'backend': backend_name.lower()}, config_file)
|
|
print('Setting the default backend to "{}". You can change it in the '
|
|
'~/.dgl/config.json file or export the DGLBACKEND environment variable. '
|
|
'Valid options are: pytorch, mxnet, tensorflow (all lowercase)'.format(
|
|
backend_name))
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("default_dir", type=str, default=os.path.join(os.path.expanduser('~'), '.dgl'))
|
|
parser.add_argument("backend", nargs=1, type=str, choices=[
|
|
'pytorch', 'tensorflow', 'mxnet'], help="Set default backend")
|
|
args = parser.parse_args()
|
|
set_default_backend(args.default_dir, args.backend[0])
|