test fix

examples/test_all.sh

@@ -2,8 +2,8 @@
 
 # 定义数组变量
 partition_params=("ours" "metis" "ldg" "random")
-partitions="4"
-nnodes="1"
+partitions="16"
+nnodes="4"
 node_rank="0"
 probability_params=("1" "0.5" "0.1" "0.05" "0.01" "0")
 sample_type_params=("recent" "boundery_recent_decay" "boundery_recent_uniform")
@@ -16,7 +16,8 @@ mkdir -p all
 
 # 遍历数组并执行命令
 for data in "${data_param[@]}"; do
-    torchrun --nnodes "$nnodes" --node_rank 0 --nproc-per-node 1 --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition ours --memory_type local --sample_type recent --topk 0 > all/"$data"-1.out &
+    mkdir all/"$data"
+    torchrun --nnodes "$nnodes" --node_rank 0 --nproc-per-node 1 --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition ours --memory_type local --sample_type recent --topk 0 > all/"$data"/1.out &
     wait
     for partition in "${partition_params[@]}"; do
         for sample in "${sample_type_params[@]}"; do
@@ -25,15 +26,15 @@ for data in "${data_param[@]}"; do
                     if [ "$mem" = "historical" ]; then
                         for ssim in "${shared_memory_ssim[@]}"; do
                             if [ "$partition" = "ours" ]; then
-                                torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0.01 --sample_type "$sample" --memory_type "$mem" --shared_memory_ssim "$ssim" > all/"$data"-"$partitions"-ours_shared-0.01-"$mem"-"$ssim"-"$sample".out &
+                                torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0.01 --sample_type "$sample" --memory_type "$mem" --shared_memory_ssim "$ssim" > all/"$data"/"$partitions"-ours_shared-0.01-"$mem"-"$ssim"-"$sample".out &
                                 wait
                             fi
                         done
                     else
-                        torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0 --sample_type "$sample" --memory_type "$mem" > all/"$data"-"$partitions"-"$partition"-0-"$mem"-"$sample".out &
+                        torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0 --sample_type "$sample" --memory_type "$mem" > all/"$data"/"$partitions"-"$partition"-0-"$mem"-"$sample".out &
                         wait
                         if [ "$partition" = "ours" ]; then
-                            torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0.01 --sample_type "$sample" --memory_type "$mem" > all/"$data"-"$partitions"-ours_shared-0.01-"$mem"-"$sample".out &
+                            torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0.01 --sample_type "$sample" --memory_type "$mem" > all/"$data"/"$partitions"-ours_shared-0.01-"$mem"-"$sample".out &
                             wait
                         fi
                     fi
@@ -44,15 +45,15 @@ for data in "${data_param[@]}"; do
                         if [ "$mem" = "historical" ]; then
                             for ssim in "${shared_memory_ssim[@]}"; do
                                 if [ "$partition" = "ours" ]; then
-                                    torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0.01 --sample_type "$sample" --probability "$pro" --memory_type "$mem" --shared_memory_ssim "$ssim" > all/"$data"-"$partitions"-ours_shared-0.01"$mem"-"$ssim"-"$sample"-"$pro".out &
+                                    torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0.01 --sample_type "$sample" --probability "$pro" --memory_type "$mem" --shared_memory_ssim "$ssim" > all/"$data"/"$partitions"-ours_shared-0.01"$mem"-"$ssim"-"$sample"-"$pro".out &
                                     wait
                                 fi
                             done
                         else
-                            torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0 --sample_type "$sample" --probability "$pro" --memory_type "$mem" > all/"$data"-"$partitions"-"$partition"-0-"$mem"-"$sample"-"$pro".out &
+                            torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0 --sample_type "$sample" --probability "$pro" --memory_type "$mem" > all/"$data"/"$partitions"-"$partition"-0-"$mem"-"$sample"-"$pro".out &
                             wait
                             if [ "$partition" = "ours" ]; then
-                                torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0.01 --sample_type "$sample" --probability "$pro" --memory_type "$mem" > all/"$data"-"$partitions"-ours_shared-0.01-"$mem"-"$sample"-"$pro".out &
+                                torchrun --nnodes "$nnodes" --node_rank "$node_rank" --nproc-per-node "$partitions" --master-addr 192.168.1.105 --master-port 9445 train_boundery.py --dataname "$data" --mode TGN_large --partition "$partition" --topk 0.01 --sample_type "$sample" --probability "$pro" --memory_type "$mem" > all/"$data"/"$partitions"-ours_shared-0.01-"$mem"-"$sample"-"$pro".out &
                                 wait
                             fi
                         fi

examples/train_boundery.py

@@ -2,6 +2,7 @@ import argparse
 import os
 import profile
 import sys
+import psutil
 from os.path import abspath, join, dirname
 current_path = os.path.dirname(os.path.abspath(__file__))
 parent_path = os.path.abspath(os.path.join(current_path, os.pardir))
@@ -83,7 +84,15 @@ import random
 import numpy as np
 from sklearn.metrics import average_precision_score, roc_auc_score
 from torch.nn.parallel import DistributedDataParallel as DDP
-
+def get_network_interfaces_with_prefix(prefixs):
+    interfaces = psutil.net_if_addrs()
+    matching_interfaces = [iface for iface in interfaces if iface.startswith(prefixs[0]) or iface.startswith(prefixs[1])]
+    return matching_interfaces
+
+# Example usage
+prefix = ("ens4f1np1","ens6f0np0")
+matching_interfaces = get_network_interfaces_with_prefix(prefix)
+print(f"Network interfaces with prefix '{prefix}': {matching_interfaces}")
 #os.environ['CUDA_VISIBLE_DEVICES'] = '2'#str(args.rank)
 if not 'WORLD_SIZE'  in os.environ:
     os.environ["RANK"] = str(args.rank)
@@ -94,7 +103,7 @@ if not 'MASTER_ADDR' in os.environ:
 if not 'MASTER_PORT' in os.environ:
     os.environ["MASTER_PORT"] = '9337'
 os.environ["NCCL_IB_DISABLE"]='1'
-os.environ['NCCL_SOCKET_IFNAME']='ens6f0np0'
+os.environ['NCCL_SOCKET_IFNAME']=matching_interfaces[0]
 torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 local_rank = int(os.environ["LOCAL_RANK"])
 def seed_everything(seed=42):
@@ -575,8 +584,8 @@ def main():
             print('\ttrain loss:{:.4f}  train ap:{:4f}  val ap:{:4f}  val auc:{:4f} test ap {:4f} test auc{:4f}\n'.format(total_loss,train_ap, ap, auc,test_ap,test_auc))
             print('\ttotal time:{:.2f}s  prep time:{:.2f}s\n'.format(time.time()-epoch_start_time, time_prep))    
             torch.save(model.module.state_dict(), get_checkpoint_path(e))
-    torch.save(val_list,'all/{}_val_{}_{}_{}_{}_{}_{}_{}_{}.pt'.format(args.dataname,args.partition,args.topk,dist.get_world_size(),dist.get_rank(),args.sample_type,args.probability,args.memory_type,args.shared_memory_ssim))
-    torch.save(loss_list,'all/{}_loss_{}_{}_{}_{}_{}_{}_{}_{}.pt'.format(args.dataname,args.partition,args.topk,dist.get_world_size(),dist.get_rank(),args.sample_type,args.probability,args.memory_type,args.shared_memory_ssim))
+    torch.save(val_list,'all/{}/val_{}_{}_{}_{}_{}_{}_{}_{}.pt'.format(args.dataname,args.partition,args.topk,dist.get_world_size(),dist.get_rank(),args.sample_type,args.probability,args.memory_type,args.shared_memory_ssim))
+    torch.save(loss_list,'all/{}/loss_{}_{}_{}_{}_{}_{}_{}_{}.pt'.format(args.dataname,args.partition,args.topk,dist.get_world_size(),dist.get_rank(),args.sample_type,args.probability,args.memory_type,args.shared_memory_ssim))
 
     print(avg_time)
     if not early_stop:        

starrygl/sample/memory/change.py

@@ -18,9 +18,9 @@ class MemoryMoniter:
         self.memory_ssim.append(self.ssim(pre_memory,now_memory,method = 'cos'))
         self.nid_list.append(nid)
     def draw(self,degree,data,e):
-        torch.save(self.nid_list,'all/{}_nid_{}.pt'.format(data,e))
-        torch.save(self.memorychange,'all/{}_memoryF_{}.pt'.format(data,e))
-        torch.save(self.memory_ssim,'all/{}_memcos_{}.pt'.format(data,e))
+        torch.save(self.nid_list,'all/{}/memorynid_{}.pt'.format(data,e))
+        torch.save(self.memorychange,'all/{}/memoryF_{}.pt'.format(data,e))
+        torch.save(self.memory_ssim,'all/{}/memcos_{}.pt'.format(data,e))
         # path = './memory/{}/'.format(data)
 
         # if not os.path.exists(path):