refactor

benchmarks/cugraph/standalone/bulk_sampling/bench_cugraph_training.py

@@ -197,7 +197,7 @@ def main(args):
     fanout = [int(f) for f in args.fanout.split("_")]
 
     if args.framework == "PyG":
-        from trainers_pyg import PyGNativeTrainer
+        from trainers.pyg import PyGNativeTrainer
         trainer = PyGNativeTrainer(
             model=args.model,
             dataset=dataset,
@@ -215,7 +215,7 @@ def main(args):
             args.sample_dir,
             f'ogbn_papers100M[{args.replication_factor}]_b{args.batch_size}_f{fanout}'
         )
-        from trainers_cugraph_pyg import PyGCuGraphTrainer
+        from trainers.pyg import PyGCuGraphTrainer
         trainer = PyGCuGraphTrainer(
             model=args.model,
             dataset=dataset,

benchmarks/cugraph/standalone/bulk_sampling/models/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

benchmarks/cugraph/standalone/bulk_sampling/models/pyg/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .models_cugraph_pyg import CuGraphSAGE
+from .models_pyg import GraphSAGE

benchmarks/cugraph/standalone/bulk_sampling/trainers/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .trainer import Trainer
+from .trainer import extend_tensor

benchmarks/cugraph/standalone/bulk_sampling/trainers/pyg/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .trainers_cugraph_pyg import PyGCuGraphTrainer
+from .trainers_pyg import PyGNativeTrainer

benchmarks/cugraph/standalone/bulk_sampling/trainers_cugraph_pyg.py → benchmarks/cugraph/standalone/bulk_sampling/trainers/pyg/trainers_cugraph_pyg.py

@@ -11,20 +11,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from trainers_pyg import PyGTrainer
-from models_cugraph_pyg import CuGraphSAGE
+from .trainers_pyg import PyGTrainer
+from models.pyg import CuGraphSAGE
 
 import torch
 import numpy as np
 
-from torch.distributed.optim import ZeroRedundancyOptimizer
 from torch.nn.parallel import DistributedDataParallel as ddp
 
 from cugraph.gnn import FeatureStore
 from cugraph_pyg.data import CuGraphStore
 from cugraph_pyg.loader import BulkSampleLoader
 
-import gc
 import os
 
 

benchmarks/cugraph/standalone/bulk_sampling/trainers_pyg.py → benchmarks/cugraph/standalone/bulk_sampling/trainers/pyg/trainers_pyg.py

@@ -12,10 +12,10 @@
 # limitations under the License.
 
 
-from trainer import Trainer
-from trainer import extend_tensor
+from trainers import Trainer
+from trainers import extend_tensor
 from datasets import OGBNPapers100MDataset
-from models_pyg import GraphSAGE
+from models.pyg import GraphSAGE
 
 import torch
 import numpy as np