Star887
1399
代码
代码
Issues1797
Pull Requests1761
Actions
项目讨论
Wiki
分析
Star887
1399
已合并
add inference script in benchmarks #34570
HandsoemLemon创建于 4月28日
add inference script in benchmarks #34570
已合并
HandsoemLemon创建于 4月28日
从HandsoemLemon/pytorch_benchmarksv2.7.1合入到Ascend/pytorchv2.7.1
共 7 个文件变更+1587-0
Abenchmarks/inference/DCNv2/dcnv2.patch+22-0
@@ -0,0 +1,22 @@
1+--- CTR_Algorithm/pytorch/model/DCNv2.py 2026-03-31 11:09:49.207077300 +0800
2++++ dcn/CTR_Algorithm/pytorch/model/DCNv2.py 2026-03-31 11:12:56.799077300 +0800
3+@@ -120,7 +120,8 @@
4+ """
5+ super(DeepCrossNetv2, self).__init__()
6+ self.feature_fields = feature_fields
7+- self.offsets = np.array((0, *np.cumsum(feature_fields)[:-1]), dtype = np.long)
8++ self.offsets = np.array((0, *np.cumsum(feature_fields)[:-1]), dtype = np.int64)
9++ self.register_buffer("offsets_tensor", torch.as_tensor(self.offsets, dtype=torch.int64))
10+ self.model_method = model_method
11+
12+ # Embedding layer
13+@@ -157,7 +158,7 @@
14+ raise NotImplementedError
15+
16+ def forward(self, x):
17+- tmp = x + x.new_tensor(self.offsets).unsqueeze(0)
18++ tmp = x + self.offsets_tensor.unsqueeze(0)
19+
20+ # embeded dense vector
21+ embeded_x = self.embedding(tmp).view(-1, self.embedding_out_dim)
22+
Abenchmarks/inference/DCNv2/dcnv2.py+205-0
@@ -0,0 +1,205 @@
1+# -*- coding: utf-8 -*-
2+import os
3+import pandas as pd
4+import numpy as np
5+import torch
6+import torch.nn as nn
7+import torch.optim as optim
8+import torch.utils.data as Data
9+import torch.nn.functional as F
10+from sklearn.preprocessing import LabelEncoder
11+from sklearn.model_selection import train_test_split
12+import sys
13+import time
14+import argparse
15+ 
16+from CTR_Algorithm.data.AmazonDataPreprocess import AmazonBookPreprocess
17+ 
18+MODEL_NAME='DCNV2'
19+ 
20+data = pd.read_csv('CTR_Algorithm/data/data.csv')
21+data_X = data.iloc[:,2:]
22+data_y = data.click.values
23+ 
24+data_X = data_X.apply(LabelEncoder().fit_transform)
25+ 
26+fields = data_X.max().values + 1
27+ 
28+# data_X = data.iloc[:,:-1]
29+# data_y = data.label.values
30+ 
31+tmp_X, test_X, tmp_y, test_y = train_test_split(data_X, data_y, test_size = 0.2, random_state=42, stratify=data_y)
32+train_X, val_X, train_y, val_y = train_test_split(tmp_X, tmp_y, test_size = 0.25, random_state=42, stratify=tmp_y)
33+ 
34+train_X = torch.from_numpy(train_X.values).long()
35+val_X = torch.from_numpy(val_X.values).long()
36+test_X = torch.from_numpy(test_X.values).long()
37+ 
38+train_y = torch.from_numpy(train_y).long()
39+val_y = torch.from_numpy(val_y).long()
40+test_y = torch.from_numpy(test_y).long()
41+ 
42+train_set = Data.TensorDataset(train_X, train_y)
43+val_set = Data.TensorDataset(val_X, val_y)
44+train_loader = Data.DataLoader(dataset=train_set,
45+ batch_size=32,
46+ shuffle=True)
47+val_loader = Data.DataLoader(dataset=val_set,
48+ batch_size=32,
49+ shuffle=False)
50+ 
51+ 
52+def detect_device_type():
53+ try:
54+ import torch
55+ if torch.cuda.is_available():
56+ return "cuda"
57+ try:
58+ import torch_npu
59+ import os
60+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
61+ if torch.npu.is_available():
62+ return "npu"
63+ except ImportError:
64+ pass
65+ except ImportError:
66+ pass
67+ return "cpu"
68+ 
69+ 
70+def get_profile(profiler_start_step: int, profiler_end_step: int, profiling_save_path: str):
71+ warm_step = profiler_start_step
72+ active_step = profiler_end_step - warm_step +1
73+ print(f"[Profile INFO]: warm_step: {warm_step}, active_step: {active_step}, profiling_save_path: {profiling_save_path}")
74+ device = detect_device_type()
75+ if device == 'npu':
76+ import torch_npu
77+ g_prof_config = torch_npu.profiler._ExperimentalConfig(
78+ export_type=[
79+ torch_npu.profiler.ExportType.Text,
80+ torch_npu.profiler.ExportType.Db
81+ ],
82+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
83+ msprof_tx=False,
84+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
85+ l2_cache=False,
86+ op_attr=False,
87+ data_simplification=False,
88+ record_op_args=False,
89+ gc_detect_threshold=None)
90+ 
91+ return torch_npu.profiler.profile(
92+ activities=[
93+ torch_npu.profiler.ProfilerActivity.CPU,
94+ torch_npu.profiler.ProfilerActivity.NPU],
95+ schedule=torch_npu.profiler.schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
96+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_save_path),
97+ record_shapes=False,
98+ profile_memory=False,
99+ with_stack=False,
100+ with_modules=False,
101+ with_flops=False,
102+ experimental_config=g_prof_config)
103+ elif device == 'cuda':
104+ from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity
105+ return profile(
106+ activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
107+ schedule=schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
108+ on_trace_ready=tensorboard_trace_handler(profiling_save_path),
109+ record_shapes=False,
110+ profile_memory=False,
111+ )
112+ else:
113+ print(f"[Warning]: No supported acceleration device (CUDA/NPU) detected, profiler will be disabled")
114+ return None
115+ 
116+def detect_device_type():
117+ try:
118+ import torch
119+ if torch.cuda.is_available():
120+ return "cuda"
121+ try:
122+ import torch_npu
123+ import os
124+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
125+ if torch.npu.is_available():
126+ return "npu"
127+ except ImportError:
128+ pass
129+ except ImportError:
130+ pass
131+ return "cpu"
132+ 
133+def device_synchronize():
134+ if torch.cuda.is_available():
135+ torch.cuda.synchronize()
136+ elif torch.npu.is_available():
137+ torch.npu.synchronize()
138+ 
139+def eval_op_prof(model: nn.Module, device, mod, args):
140+ dl = iter(val_loader)
141+ x, _ = next(dl)
142+ x = x.to(device)
143+ pt_path= MODEL_NAME +'_'+ mod + '_result.pt'
144+ 
145+ model.eval()
146+ with torch.no_grad():
147+ res = model(x)
148+ torch.save(res, pt_path)
149+ device_synchronize()
150+ 
151+ prof = None
152+ if args.enable_profiler:
153+ profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mod
154+ prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
155+ prof.start()
156+ 
157+ execution_times = []
158+ with torch.no_grad():
159+ for i in range(args.max_steps):
160+ start_time = time.time()
161+ model(x)
162+ device_synchronize()
163+ end_time = time.time()
164+ step_time_ms = (end_time - start_time) * 1000
165+ print(f"[{mod}]: Step {i}: {step_time_ms:.4f} ms")
166+
167+ if i >= 10:
168+ execution_times.append(step_time_ms)
169+ 
170+ if args.enable_profiler:
171+ prof.step()
172+
173+ if args.enable_profiler:
174+ prof.stop()
175+ if execution_times:
176+ avg_ms = sum(execution_times) / len(execution_times)
177+ print(f"[{mod}]:: Avg over {len(execution_times)} steps: {avg_ms:.4f} ms")
178+ 
179+parser = argparse.ArgumentParser(description=MODEL_NAME + " infernece")
180+parser.add_argument("--max_steps", type=int, default=200,
181+ help="Total training steps")
182+parser.add_argument("--enable_compile", action="store_true",
183+ help="Enable torch.compile and Inductor backend")
184+parser.add_argument("--enable_profiler", action="store_true",
185+ help="Enable profiler for performance analysis")
186+parser.add_argument("--profiler_start_step", type=int, default=5,
187+ help="Output directory for trained model")
188+parser.add_argument("--profiler_end_step", type=int, default=8,
189+ help="Output directory for trained model")
190+parser.add_argument("--profiler_save_path", type=str, default="./profile",
191+ help="Output directory for trained model")
192+args = parser.parse_args()
193+ 
194+from CTR_Algorithm.pytorch.model.DCNv2 import DeepCrossNetv2
195+model = DeepCrossNetv2(feature_fields = fields, embed_dim = 16, layer_num = 2, mlp_dims = (32, 16), dropout = 0.1, cross_method='Matrix')
196+ 
197+device_type = detect_device_type()
198+device = torch.device(device_type)
199+model.to(device)
200+mod = 'eager'
201+if args.enable_compile:
202+ model = torch.compile(model, dynamic=False)
203+ mod = 'compile'
204+print(f"{MODEL_NAME} inference begin, mode is {mod}")
205+eval_op_prof(model, device, mod, args)
Abenchmarks/inference/DIN/din.py+176-0
@@ -0,0 +1,176 @@
1+# -*- coding: utf-8 -*-
2+import os
3+import time
4+import argparse
5+import sys
6+import pandas as pd
7+import numpy as np
8+import torch
9+import torch.nn as nn
10+import torch.optim as optim
11+import torch.utils.data as Data
12+import torch.nn.functional as F
13+from sklearn.model_selection import train_test_split
14+ 
15+from CTR_Algorithm.pytorch.model import DIN
16+from CTR_Algorithm.data.AmazonDataPreprocess import AmazonBookPreprocess
17+ 
18+MODEL_NAME = "DIN"
19+ 
20+data = pd.read_csv('CTR_Algorithm/data/amazon-books-100k.txt')
21+data = AmazonBookPreprocess(data)
22+fields = data.max().max()
23+ 
24+data_X = data.iloc[:,:-1]
25+data_y = data.label.values
26+ 
27+tmp_X, test_X, tmp_y, test_y = train_test_split(data_X, data_y, test_size = 0.2, random_state=42, stratify=data_y)
28+train_X, val_X, train_y, val_y = train_test_split(tmp_X, tmp_y, test_size = 0.25, random_state=42, stratify=tmp_y)
29+ 
30+train_X = torch.from_numpy(train_X.values).long()
31+val_X = torch.from_numpy(val_X.values).long()
32+test_X = torch.from_numpy(test_X.values).long()
33+ 
34+train_y = torch.from_numpy(train_y).long()
35+val_y = torch.from_numpy(val_y).long()
36+test_y = torch.from_numpy(test_y).long()
37+ 
38+train_set = Data.TensorDataset(train_X, train_y)
39+val_set = Data.TensorDataset(val_X, val_y)
40+train_loader = Data.DataLoader(dataset=train_set, batch_size=32, shuffle=True)
41+val_loader = Data.DataLoader(dataset=val_set, batch_size=32, shuffle=False)
42+ 
43+def get_profile(profiler_start_step: int, profiler_end_step: int, profiling_save_path: str):
44+ warm_step = profiler_start_step
45+ active_step = profiler_end_step - warm_step +1
46+ print(f"[Profile INFO]: warm_step: {warm_step}, active_step: {active_step}, profiling_save_path: {profiling_save_path}")
47+ device = detect_device_type()
48+ if device == 'npu':
49+ import torch_npu
50+ g_prof_config = torch_npu.profiler._ExperimentalConfig(
51+ export_type=[
52+ torch_npu.profiler.ExportType.Text,
53+ torch_npu.profiler.ExportType.Db
54+ ],
55+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
56+ msprof_tx=False,
57+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
58+ l2_cache=False,
59+ op_attr=False,
60+ data_simplification=False,
61+ record_op_args=False,
62+ gc_detect_threshold=None)
63+ 
64+ return torch_npu.profiler.profile(
65+ activities=[
66+ torch_npu.profiler.ProfilerActivity.CPU,
67+ torch_npu.profiler.ProfilerActivity.NPU],
68+ schedule=torch_npu.profiler.schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
69+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_save_path),
70+ record_shapes=False,
71+ profile_memory=False,
72+ with_stack=False,
73+ with_modules=False,
74+ with_flops=False,
75+ experimental_config=g_prof_config)
76+ elif device == 'cuda':
77+ from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity
78+ return profile(
79+ activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
80+ schedule=schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
81+ on_trace_ready=tensorboard_trace_handler(profiling_save_path),
82+ record_shapes=False,
83+ profile_memory=False,
84+ )
85+ else:
86+ print(f"[Warning]: No supported acceleration device (CUDA/NPU) detected, profiler will be disabled")
87+ return None
88+ 
89+def detect_device_type():
90+ try:
91+ import torch
92+ if torch.cuda.is_available():
93+ return "cuda"
94+ try:
95+ import torch_npu
96+ import os
97+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
98+ if torch.npu.is_available():
99+ return "npu"
100+ except ImportError:
101+ pass
102+ except ImportError:
103+ pass
104+ return "cpu"
105+ 
106+def device_synchronize():
107+ if torch.cuda.is_available():
108+ torch.cuda.synchronize()
109+ elif torch.npu.is_available():
110+ torch.npu.synchronize()
111+ 
112+def eval_op_prof(model: nn.Module, device, mod, args):
113+ dl = iter(val_loader)
114+ x, _ = next(dl)
115+ x = x.to(device)
116+ pt_path= MODEL_NAME +'_'+ mod + '_result.pt'
117+ 
118+ model.eval()
119+ with torch.no_grad():
120+ res = model(x)
121+ torch.save(res, pt_path)
122+ device_synchronize()
123+ 
124+ prof = None
125+ if args.enable_profiler:
126+ profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mod
127+ prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
128+ prof.start()
129+ 
130+ execution_times = []
131+ with torch.no_grad():
132+ for i in range(args.max_steps):
133+ start_time = time.time()
134+ model(x)
135+ device_synchronize()
136+ end_time = time.time()
137+ step_time_ms = (end_time - start_time) * 1000
138+ print(f"[{mod}]: Step {i}: {step_time_ms:.4f} ms")
139+
140+ if i >= 10:
141+ execution_times.append(step_time_ms)
142+ 
143+ if args.enable_profiler:
144+ prof.step()
145+
146+ if args.enable_profiler:
147+ prof.stop()
148+ if execution_times:
149+ avg_ms = sum(execution_times) / len(execution_times)
150+ print(f"[{mod}]:: Avg over {len(execution_times)} steps: {avg_ms:.4f} ms")
151+ 
152+parser = argparse.ArgumentParser(description=MODEL_NAME + " infernece")
153+parser.add_argument("--max_steps", type=int, default=200,
154+ help="Total training steps")
155+parser.add_argument("--enable_compile", action="store_true",
156+ help="Enable torch.compile and Inductor backend")
157+parser.add_argument("--enable_profiler", action="store_true",
158+ help="Enable profiler for performance analysis")
159+parser.add_argument("--profiler_start_step", type=int, default=5,
160+ help="Output directory for trained model")
161+parser.add_argument("--profiler_end_step", type=int, default=8,
162+ help="Output directory for trained model")
163+parser.add_argument("--profiler_save_path", type=str, default="./profile",
164+ help="Output directory for trained model")
165+args = parser.parse_args()
166+ 
167+model = DIN.DeepInterestNet(feature_dim=fields, embed_dim=8, mlp_dims=[64,32], dropout=float(0.2))
168+device_type = detect_device_type()
169+device = torch.device(device_type)
170+model.to(device)
171+mod = 'eager'
172+if args.enable_compile:
173+ model = torch.compile(model, dynamic=False)
174+ mod = 'compile'
175+print(f"{MODEL_NAME} inference begin, mode is {mod}")
176+eval_op_prof(model, device, mod, args)
Abenchmarks/inference/ETA/eta.py+510-0
@@ -0,0 +1,510 @@
1+import argparse
2+import glob
3+import json
4+import os
5+import random
6+import stat
7+import time
8+from datetime import datetime
9+ 
10+import h5py
11+import numpy as np
12+ 
13+import torch
14+import torch.nn as nn
15+import torch.optim as optim
16+import torch.nn.functional as F
17+from torch.nn.utils.rnn import pad_sequence
18+from torch.utils.data import Dataset, DataLoader, ConcatDataset
19+from sklearn.metrics import roc_auc_score
20+ 
21+torch.manual_seed(2024)
22+random.seed(2024)
23+ 
24+MODEL_NAME = "ETA"
25+STD_DEV = (2 / 512) ** 0.5
26+EMBEDDING_FEATURE_NUM = 27
27+TARGET_FIELDS = ["206", "207", "216", "210"]
28+ 
29+def detect_device_type():
30+ try:
31+ import torch
32+ if torch.cuda.is_available():
33+ return "cuda"
34+ try:
35+ import torch_npu
36+ import os
37+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
38+ if torch.npu.is_available():
39+ return "npu"
40+ except ImportError:
41+ pass
42+ except ImportError:
43+ pass
44+ return "cpu"
45+ 
46+def json_file_load(json_name: str, json_path: str) -> dict:
47+ """
48+ Load a JSON file from the specified path.
49+ """
50+ flags = os.O_RDONLY
51+ modes = stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH
52+ try:
53+ with os.fdopen(os.open(json_path, flags, modes), "r") as fp:
54+ json_re = json.load(fp)
55+ except FileNotFoundError as e:
56+ raise FileNotFoundError(f"{json_name} file not found: {e}") from e
57+ except Exception as e:
58+ raise RuntimeError(f"Error loading {json_name} file: {e}") from e
59+ 
60+ return json_re
61+ 
62+def collate_fn(batch):
63+ input_dicts = [item[0] for item in batch]
64+ target_dicts = [item[1] for item in batch]
65+ input_tensors = {}
66+ for key in input_dicts[0].keys():
67+ tensors = [d[key] for d in input_dicts if key in d and d[key] is not None]
68+ if not tensors:
69+ continue
70+ if tensors[0].dim() == 0:
71+ tensors = [t.unsqueeze(0) for t in tensors]
72+ input_tensors[key] = pad_sequence(tensors, batch_first=True)
73+ 
74+ target_tensors = {}
75+ for key in target_dicts[0].keys():
76+ tensors = [d[key] for d in target_dicts if key in d and d[key] is not None]
77+ if not tensors:
78+ continue
79+ target_tensors[key] = torch.stack(tensors)
80+ 
81+ return input_tensors, target_tensors
82+ 
83+def parse_arguments():
84+ parser = argparse.ArgumentParser()
85+ parser.add_argument("--embedding_size", type=int, default=16, help="Embedding size")
86+ parser.add_argument("--attention_dim", type=int, default=64, help="")
87+ parser.add_argument("--enable_compile", action="store_true", help="")
88+ parser.add_argument("--num_heads", type=int, default=4, help="")
89+ parser.add_argument("--short_output_dim", type=int, default=16, help="short attention output dimension")
90+ parser.add_argument("--max_seq_len", type=int, default=50, help="")
91+ parser.add_argument("--topk", type=int, default=16, help="")
92+ parser.add_argument("--deep_layers", type=str, default="512,256,128,64", help="deep layers")
93+ parser.add_argument("--hash_bits", type=int, default=32, help="")
94+ parser.add_argument("--reuse_hash", type=bool, default=True, help="")
95+ parser.add_argument("--batch_size", type=int, default=4096, help="Number of batch size")
96+ parser.add_argument("--learning_rate", type=float, default=0.001, help="learning rate")
97+ parser.add_argument("--optimizer", type=str, default="Adam",
98+ choices=["Adam", "Adagrad", "GD", "Momentum"], help="")
99+ parser.add_argument('--early_stop_patience', type=int, default=5, help="")
100+ parser.add_argument("--data_dir", type=str, default="./aliccp/aliccp_out/", help="data dir")
101+ parser.add_argument("--dt_dir", type=str, default='', help="data dt partition")
102+ parser.add_argument("--model_dir", type=str, default=f"./", help="code check point dir")
103+ parser.add_argument("--clear_existing_model", action="store_true", help="")
104+ parser.add_argument("--task_type", type=str, default="train",
105+ choices=["train", "eval", "predict"], help="task type")
106+ parser.add_argument("--max_steps", type=int, default=100,
107+ help="Total training steps (优先级高于num_epochs，-1表示使用num_epochs控制)")
108+ parser.add_argument('--epoch_num', type=int, default=1, help="Number of epochs")
109+ parser.add_argument('--train_batch_num', type=int, default=2000, help="Number of train batchs")
110+ parser.add_argument('--test_batch_num', type=int, default=10, help="Number of test batchs")
111+ parser.add_argument("--enable_profiler", action="store_true",
112+ help="Enable profiler for performance analysis")
113+ parser.add_argument("--profiler_start_step", type=int, default=5,
114+ help="Output directory for trained model")
115+ parser.add_argument("--profiler_end_step", type=int, default=8,
116+ help="Output directory for trained model")
117+ parser.add_argument("--profiler_save_path", type=str, default="./profile",
118+ help="Output directory for trained model")
119+ return parser.parse_args()
120+ 
121+ 
122+class HDF5Dataset(Dataset):
123+ def __init__(self, hdf5_path):
124+ self.hdf5_path = hdf5_path
125+ self._load_hdf5()
126+ 
127+ def _load_hdf5(self):
128+ with h5py.File(self.hdf5_path, 'r') as f:
129+ self.input_sample = {}
130+ self.target_sample = {}
131+ y = np.array(f["y"])
132+ z = np.array(f["z"])
133+ fields = [key for key in f.keys() if key not in ["y", "z"]]
134+ self.target_sample.update({"y": torch.tensor(y, dtype=torch.float32)})
135+ self.target_sample.update({"z": torch.tensor(z, dtype=torch.float32)})
136+ for multi_field in fields:
137+ self.input_sample.update({multi_field: torch.tensor(np.array(f[multi_field]), dtype=torch.int64)})
138+ self._length = len(y)
139+ 
140+ def __len__(self):
141+ return self._length
142+ 
143+ def __getitem__(self, idx):
144+ input_dict = {k: v[idx] for k, v in self.input_sample.items()}
145+ target_dict = {k: v[idx] for k, v in self.target_sample.items()}
146+ return input_dict, target_dict
147+ 
148+ 
149+class TorchDataSet(ConcatDataset):
150+ def __init__(self, files):
151+ datasets = [HDF5Dataset(fp) for fp in files]
152+ super().__init__(datasets)
153+ 
154+ 
155+class ETA(nn.Module):
156+ def __init__(self, spec, params):
157+ super(ETA, self).__init__()
158+ self.spec = spec
159+ self.params = params
160+ 
161+ # embedding layers
162+ self.emb_weights = nn.ModuleDict()
163+ for key, vocab_len in spec["vocab_length"].items():
164+ self.emb_weights[key] = nn.Embedding(vocab_len + 1, params.embedding_size)
165+ nn.init.normal_(self.emb_weights[key].weight, std=STD_DEV)
166+ 
167+ self.hash_weights = nn.Parameter(
168+ torch.randn(params.embedding_size, params.hash_bits),
169+ requires_grad=False
170+ )
171+ 
172+ # attention paramters
173+ self.short_attentions = nn.ModuleList()
174+ self.long_attentions = nn.ModuleList()
175+ for _ in range(len(TARGET_FIELDS)): # 4 target fields
176+ self.short_attentions.append(ShortAttention(params))
177+ self.long_attentions.append(LongAttention(self.hash_weights, params))
178+ 
179+ # mlp layer
180+ self.mlp = nn.Sequential()
181+ deep_layers = list(map(int, params.deep_layers.strip().split(",")))
182+ input_dim = EMBEDDING_FEATURE_NUM * params.embedding_size
183+ for i, dim in enumerate(deep_layers):
184+ self.mlp.add_module(f'mlp{i}', nn.Linear(input_dim, dim))
185+ self.mlp.add_module(f'relu{i}', nn.ReLU())
186+ input_dim = dim
187+ 
188+ # output layer
189+ self.output = nn.Linear(input_dim, 1)
190+ 
191+ def embedding_lookup_sparse(self, params: nn.Embedding, ids: torch.Tensor, combiner: str):
192+ mask = (ids >= 0).float().unsqueeze(-1)
193+ ids = ids.clone()
194+ ids[ids == -1] = 0
195+ embedding = params(ids) * mask
196+ if combiner == "sum":
197+ return embedding.sum(dim=1)
198+ elif combiner == "mean":
199+ return embedding.sum(dim=1) / mask.sum(dim=1).clamp(min=1e-7)
200+ else:
201+ raise ValueError("combiner only supoort 'sum', 'mean'")
202+ 
203+ def forward(self, features):
204+ # embedding
205+ embeddings = {}
206+ masks = {}
207+ for key in self.spec["one_hot_fields"]:
208+ embeddings[key] = self.emb_weights[key](features[key]).unsqueeze(1)
209+ for key in self.spec["multi_hot_fields"]:
210+ feat = features[key]
211+ masks[key] = (feat >= 0).bool().unsqueeze(1)
212+ feat = feat.clone()
213+ feat[feat == -1] = 0
214+ embeddings[key] = self.emb_weights[key](feat)
215+ for key in self.spec["special_fields"]:
216+ embeddings[key] = self.embedding_lookup_sparse(
217+ self.emb_weights[key], features[key], combiner="sum"
218+ ).unsqueeze(1)
219+ 
220+ def long_emb_cat(field_name):
221+ dense_embedding = embeddings.get(field_name)
222+ dense_mask = masks.get(field_name)
223+ if dense_embedding is None or dense_mask is None:
224+ raise ValueError(f"Field {field_name} not found in embeddings or masks")
225+ padded_embedding = F.pad(dense_embedding, (0, 0, 0, self.params.max_seq_len), "constant", 0)
226+ padded_mask = F.pad(dense_mask, (0, self.params.max_seq_len), "constant", False)
227+ return (
228+ padded_embedding[:, :self.params.topk, :],
229+ padded_embedding[:, :self.params.max_seq_len, :],
230+ padded_mask[:, :, :self.params.topk],
231+ padded_mask[:, :, :self.params.max_seq_len]
232+ )
233+ emb_cats = [long_emb_cat(field) for field in self.spec["multi_hot_fields"]]
234+ target_fields = TARGET_FIELDS
235+ 
236+ short_attns = []
237+ long_attns = []
238+ 
239+ for i, (emb_cat, target) in enumerate(zip(emb_cats, target_fields)):
240+ emb_target = embeddings[target]
241+ 
242+ emb_short = emb_cat[0]
243+ mask_short = emb_cat[2]
244+ short_attns.append(self.short_attentions[i](emb_target, emb_short, mask_short))
245+ 
246+ emb_long = emb_cat[1]
247+ mask_long = emb_cat[3]
248+ long_attns.append(self.long_attentions[i](emb_target, emb_long, mask_long))
249+ 
250+ # concat all embeddings
251+ all_embs = []
252+ for field in self.spec["one_hot_fields"] + self.spec["special_fields"]:
253+ if embeddings[field].dim() > 3:
254+ all_embs.append(embeddings[field].squeeze(1))
255+ else:
256+ all_embs.append(embeddings[field])
257+ all_embs += short_attns + long_attns
258+ all_embs = torch.cat(all_embs, dim=1)
259+ if all_embs.dim() >= 2:
260+ concat_emb = all_embs.squeeze(1)
261+ else:
262+ concat_emb = all_embs
263+ 
264+ concat_emb = concat_emb.reshape(concat_emb.size(0), -1)
265+ # mlp and outputs
266+ mlp_out = self.mlp(concat_emb)
267+ logits = self.output(mlp_out).squeeze()
268+ pred = torch.sigmoid(logits)
269+ return pred, logits
270+ 
271+ def build_loss(self, pred, labels, click_weight=0.14, epsilon=1e-7):
272+ if pred.shape != labels.shape:
273+ raise ValueError(f"pred and labels must be the same shape. "
274+ f"pred shape: {pred.shape}, labels shape: {labels.shape}")
275+ pred = torch.clamp(pred, min=epsilon, max=1 - epsilon)
276+ loss = - (1 - click_weight) / click_weight * labels * torch.log(pred) - (1 - labels) * torch.log(1 - pred)
277+ return loss.mean()
278+ 
279+ def build_optimizer(self):
280+ if self.params.optimizer == "Adam":
281+ optimizer = optim.Adam(
282+ params=self.parameters(),
283+ lr=self.params.learning_rate,
284+ betas=[0.9, 0.999], eps=1e-8
285+ )
286+ elif self.params.optimizer == "Adagrad":
287+ optimizer = optim.Adagrad(
288+ params=self.parameters(),
289+ lr=self.params.learning_rate,
290+ initial_accumulator_value=1e-6
291+ )
292+ elif self.params.optimizer == "Momentum":
293+ optimizer = optim.SGD(
294+ params=self.parameters(),
295+ lr=self.params.learning_rate,
296+ momentum=0.95
297+ )
298+ elif self.params.optimizer == "SGD":
299+ optimizer = optim.SGD(
300+ params=self.parameters(),
301+ lr=self.params.learning_rate,
302+ )
303+ else:
304+ raise ValueError("Unsupported optimizer type: {}".format(args.optimizer))
305+ return optimizer
306+ 
307+ 
308+class ShortAttention(nn.Module):
309+ def __init__(self, params):
310+ super().__init__()
311+ self.attention_dim = params.attention_dim
312+ self.num_heads = params.num_heads
313+ self.key_dim = self.attention_dim // self.num_heads
314+ 
315+ self.q_fc = nn.Linear(params.embedding_size, self.attention_dim)
316+ self.k_fc = nn.Linear(params.embedding_size, self.attention_dim)
317+ self.v_fc = nn.Linear(params.embedding_size, self.attention_dim)
318+ self.o_fc = nn.Linear(self.attention_dim, params.short_output_dim)
319+ 
320+ def forward(self, target_input, seq_input, mask):
321+ # target_input:[B, 1, E], seq_input:[B, S, E], mask:[B, 1, S] B = target_input.size(0)
322+ # logger.info(f"tgt shape: {target_input.shape}, seq shape: {seq_input.shape}, mask shape: {mask.shape}")
323+ b, tgt_len = target_input.shape[:2]
324+ b, seq_len = seq_input.shape[:2]
325+ query = self.q_fc(target_input) # [B, 1, A]
326+ key = self.k_fc(seq_input) # [B, S, A]
327+ value = self.v_fc(seq_input) # [B, S, A]
328+ 
329+ # split heads
330+ query = query.view(b, tgt_len, self.num_heads, self.key_dim).permute(0, 2, 1, 3) # [B, Heads, 1, key_dim]
331+ key = key.view(b, seq_len, self.num_heads, self.key_dim).permute(0, 2, 3, 1) # [B, Heads, key_dim, S]
332+ value = value.view(b, seq_len, self.num_heads, self.key_dim).permute(0, 2, 1, 3) # [B, Heads, S, key_dim]
333+ 
334+ # scaled dot-product attention
335+ scores = torch.matmul(query, key) / (self.key_dim ** 0.5) # [B, Heads, 1, S]
336+ scores = scores.masked_fill(~mask.unsqueeze(1), float('-inf'))
337+ 
338+ attn = F.softmax(scores, dim=-1)
339+ 
340+ output = torch.matmul(attn, value) # [B, Heads, 1, key_dim]
341+ output = output.permute(0, 2, 1, 3)
342+ b, lens, heads, dims = output.shape
343+ output = output.reshape(b * lens, 1, heads * dims)
344+ output = self.o_fc(output)
345+ return output
346+ 
347+ 
348+class LongAttention(nn.Module):
349+ def __init__(self, hash_weights, params):
350+ super().__init__()
351+ self.hash_weights = hash_weights
352+ self.reuse_hash = params.reuse_hash
353+ self.topk = params.topk
354+ self.short_attn = ShortAttention(params)
355+ 
356+ def lsh_hash(self, vecs, hash_weights):
357+ rotated_vecs = torch.matmul(vecs, hash_weights)
358+ return (rotated_vecs > 0).float()
359+ 
360+ def forward(self, target_input, seq_input, mask):
361+ # target_input:[B, 1, E], seq_input:[B, S, E], mask:[B, 1, S]
362+ b, s, e = seq_input.shape
363+ hash_weights = self.hash_weights if self.reuse_hash else nn.Parameter(
364+ torch.randn(e, self.hash_weights.size(-1)),
365+ requires_grad=False
366+ ) # [E, hash_bits]
367+ 
368+ target_hash = self.lsh_hash(target_input.squeeze(), hash_weights) # [B, hash_bits]
369+ seq_hash = self.lsh_hash(seq_input, hash_weights) # [B, S, hash_bits]
370+ # calculate similarity(hamming distance)
371+ hash_sim = - torch.sum(torch.abs(seq_hash - target_hash.unsqueeze(1)), dim=-1) # [B, S]
372+ hash_sim = hash_sim.masked_fill(~mask.squeeze(1), float('-inf'))
373+ 
374+ # select topk
375+ _, topk_idx = torch.topk(hash_sim, self.topk, dim=-1) # [B, topk]
376+ 
377+ topk_seq = torch.gather(
378+ seq_input,
379+ 1,
380+ topk_idx.unsqueeze(-1).expand(-1, -1, e)
381+ )
382+ 
383+ topk_idx = topk_idx.unsqueeze(1)
384+ topk_mask = torch.gather(mask, 2, topk_idx)
385+ 
386+ # apply short attention to topk
387+ output = self.short_attn(target_input, topk_seq, topk_mask)
388+ return output
389+ 
390+def get_profile(profiler_start_step: int, profiler_end_step: int, profiling_save_path: str):
391+ warm_step = profiler_start_step
392+ active_step = profiler_end_step - warm_step +1
393+ print(f"[Profile INFO]: warm_step: {warm_step}, active_step: {active_step}, profiling_save_path: {profiling_save_path}")
394+ device = detect_device_type()
395+ if device == 'npu':
396+ import torch_npu
397+ g_prof_config = torch_npu.profiler._ExperimentalConfig(
398+ export_type=[
399+ torch_npu.profiler.ExportType.Text,
400+ torch_npu.profiler.ExportType.Db
401+ ],
402+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
403+ msprof_tx=False,
404+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
405+ l2_cache=False,
406+ op_attr=False,
407+ data_simplification=False,
408+ record_op_args=False,
409+ gc_detect_threshold=None)
410+ 
411+ return torch_npu.profiler.profile(
412+ activities=[
413+ torch_npu.profiler.ProfilerActivity.CPU,
414+ torch_npu.profiler.ProfilerActivity.NPU],
415+ schedule=torch_npu.profiler.schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
416+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_save_path),
417+ record_shapes=False,
418+ profile_memory=False,
419+ with_stack=False,
420+ with_modules=False,
421+ with_flops=False,
422+ experimental_config=g_prof_config)
423+ elif device == 'cuda':
424+ from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity
425+ return profile(
426+ activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
427+ schedule=schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
428+ on_trace_ready=tensorboard_trace_handler(profiling_save_path),
429+ record_shapes=False,
430+ profile_memory=False,
431+ )
432+ else:
433+ print(f"[Warning]: No supported acceleration device (CUDA/NPU) detected, profiler will be disabled")
434+ return None
435+ 
436+ 
437+def device_synchronize():
438+ if torch.cuda.is_available():
439+ torch.cuda.synchronize()
440+ elif torch.npu.is_available():
441+ torch.npu.synchronize()
442+ 
443+ 
444+def evaluate_op(model: ETA, dataloader, device, args):
445+ mod = 'compile' if args.enable_compile else 'eager'
446+ dl = iter(dataloader)
447+ features, _ = next(dl)
448+ features = {k: v.to(device) for k, v in features.items()}
449+ 
450+ for key, value in features.items():
451+ if isinstance(value, torch.Tensor) and value.dtype == torch.int64:
452+ features[key] = value.to(torch.int32)
453+ 
454+ if mod == 'compile':
455+ model = torch.compile(model, dynamic=False)
456+ model.eval()
457+ with torch.no_grad():
458+ model(features)
459+ device_synchronize()
460+ 
461+ prof = None
462+ if args.enable_profiler:
463+ profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mod
464+ prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
465+ prof.start()
466+ 
467+ exec_times = []
468+ with torch.no_grad():
469+ for i in range(args.max_steps):
470+ e2etime1 = time.perf_counter()
471+ model(features)
472+ device_synchronize()
473+ e2etime2 = time.perf_counter()
474+ elapsed_ms = (e2etime2 - e2etime1) * 1000
475+ print(f"iterations {i} / {args.max_steps}: [{mod}] e2e time: {elapsed_ms} ms.")
476+ 
477+ if i >= 10:
478+ exec_times.append(elapsed_ms)
479+ 
480+ if args.enable_profiler:
481+ prof.step()
482+ if args.enable_profiler:
483+ prof.stop()
484+ if exec_times:
485+ avg_time = sum(exec_times) / len(exec_times)
486+ print("Step time consumption statistics (excluding the first 10 steps)")
487+ print(f"\n>>> [{mod}] avg e2e: {avg_time:.3f} ms <<<")
488+ 
489+def main(args):
490+ test_files = glob.glob("%stest/data_test.csv.hd5.*" % args.data_dir)
491+ spec = json_file_load("spec", os.path.join(args.data_dir, "spec.json"))
492+ model = ETA(spec, args)
493+ device_type = detect_device_type()
494+ print(f"Using device type {device_type}")
495+ device = torch.device(device_type)
496+ model.to(device)
497+ test_dataset = TorchDataSet(test_files[1:2])
498+ test_loader = DataLoader(dataset=test_dataset,
499+ batch_size=args.batch_size,
500+ shuffle=True,
501+ collate_fn=collate_fn,
502+ prefetch_factor=100,
503+ num_workers=10)
504+ evaluate_op(model, test_loader, device, args)
505+ 
506+ 
507+if __name__ == "__main__":
508+ print(f"{MODEL_NAME} inference begin !!!")
509+ args = parse_arguments()
510+ main(args)
Abenchmarks/inference/MMOE/mmoe.py+577-0
@@ -0,0 +1,577 @@
1+import os
2+import stat
3+import glob
4+import json
5+import random
6+import shutil
7+import logging
8+import time
9+from datetime import datetime
10+import argparse
11+ 
12+import pytz
13+import h5py
14+import numpy as np
15+import torch
16+import torch.nn as nn
17+import torch.optim as optim
18+from torch.nn.utils.rnn import pad_sequence
19+from torch.utils.data import DataLoader, Dataset, ConcatDataset
20+from sklearn.metrics import roc_auc_score
21+ 
22+torch.manual_seed(2024)
23+random.seed(2024)
24+ 
25+MODEL_NAME = "MMOE"
26+EMBEDDING_FEATURE_NUM = 23
27+STD_DEV = (2 / 512) ** 0.5
28+ 
29+def detect_device_type():
30+ try:
31+ import torch
32+ if torch.cuda.is_available():
33+ return "cuda"
34+ try:
35+ import torch_npu
36+ import os
37+ os.environ['TORCHINDUCTOR_NPU_BACKEND'] = 'mlir'
38+ if torch.npu.is_available():
39+ return "npu"
40+ except ImportError:
41+ pass
42+ except ImportError:
43+ pass
44+ return "cpu"
45+ 
46+def parse_arguments():
47+ parser = argparse.ArgumentParser(description='PyTorch Example with Command Line Arguments')
48+ parser.add_argument('--embedding_size', type=int, default=16, help="Embedding size")
49+ parser.add_argument('--batch_size', type=int, default=4096, help="Batch size for training")
50+ parser.add_argument('--learning_rate', type=float, default=0.001, help="Learning rate")
51+ parser.add_argument('--optimizer', type=str, default="Adam", choices=["Adam", "Adagrad", "GD", "Momentum"],
52+ help="Optimizer type")
53+ parser.add_argument('--expert_layers', type=str, default="512,256", help="Expert layers")
54+ parser.add_argument('--tower_layers', type=str, default="128,64", help="tower layers")
55+ parser.add_argument('--ctr_task_wgt', type=float, default=0.5, help="loss weight of ctr task")
56+ parser.add_argument('--data_dir', type=str, default="./aliccp/aliccp_out/", help="Data directory")
57+ parser.add_argument('--dt_dir', type=str, default="", help="Data dt partition")
58+ parser.add_argument('--model_dir', type=str, default=f"./",
59+ help="Model checkpoint directory")
60+ parser.add_argument('--servable_model_dir', type=str, default=f"./",
61+ help="Export servable model for pytorch Serving")
62+ parser.add_argument('--clear_existing_model', action="store_true", help="Clear existing model or not")
63+ parser.add_argument('--max_seq_len', type=int, default=50, help="Max length of sequence")
64+ parser.add_argument('--task_num', type=int, default=2, help="Task number")
65+ parser.add_argument('--experts_num', type=int, default=8, help="Number of experts")
66+ parser.add_argument('--log_level', type=str, default="DEBUG",
67+ choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
68+ help="Log level")
69+ parser.add_argument('--epoch_num', type=int, default=10, help="Number of epochs")
70+ parser.add_argument('--train_batch_num', type=int, default=2000, help="Number of train batchs")
71+ parser.add_argument('--eval_batch_num', type=int, default=20, help="Number of eval batchs")
72+ parser.add_argument('--test_batch_num', type=int, default=20, help="Number of test batchs")
73+ parser.add_argument("--enable_compile", action="store_true", help="")
74+ parser.add_argument("--max_steps", type=int, default=100, help="Total training steps")
75+ parser.add_argument("--enable_profiler", action="store_true",
76+ help="Enable profiler for performance analysis")
77+ parser.add_argument("--profiler_start_step", type=int, default=5,
78+ help="Output directory for trained model")
79+ parser.add_argument("--profiler_end_step", type=int, default=8,
80+ help="Output directory for trained model")
81+ parser.add_argument("--profiler_save_path", type=str, default="./profile",
82+ help="Output directory for trained model")
83+ return parser.parse_args()
84+ 
85+ 
86+def json_file_load(json_name: str, json_path: str) -> dict:
87+ """
88+ Load a JSON file from the specified path.
89+ """
90+ flags = os.O_RDONLY
91+ modes = stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH
92+ try:
93+ with os.fdopen(os.open(json_path, flags, modes), "r") as fp:
94+ json_re = json.load(fp)
95+ except FileNotFoundError as e:
96+ raise FileNotFoundError(f"{json_name} file not found: {e}") from e
97+ except Exception as e:
98+ raise RuntimeError(f"Error loading {json_name} file: {e}") from e
99+ 
100+ return json_re
101+ 
102+ 
103+def pre_deal_dataset(dataset):
104+ num_pos = 0
105+ num_neg = 0
106+ num_pos_y = 0
107+ num_neg_y = 0
108+ dataset_index = 0
109+ for sub_ds in dataset.datasets:
110+ z = sub_ds.target_sample["z"]
111+ num_pos += (z == 1).sum().item()
112+ num_neg += (z == 0).sum().item()
113+ 
114+ y = sub_ds.target_sample["y"]
115+ num_pos_y += (y == 1).sum().item()
116+ num_neg_y += (y == 0).sum().item()
117+ 
118+ # logger.info("find pos index %s", dataset_index)
119+ dataset_index += 1
120+ return np.log(num_neg_y / num_pos_y), np.log(num_neg / num_pos)
121+ 
122+ 
123+def collate_fn(batch):
124+ input_dicts = [item[0] for item in batch]
125+ target_dicts = [item[1] for item in batch]
126+ input_tensors = {}
127+ 
128+ for key in input_dicts[0].keys():
129+ tensors = [d[key] for d in input_dicts if key in d and d[key] is not None]
130+ if not tensors:
131+ continue
132+ 
133+ if tensors[0].dim() == 0:
134+ tensors = [t.unsqueeze(0) for t in tensors]
135+ input_tensors[key] = pad_sequence(tensors, batch_first=True)
136+ target_tensors = {}
137+ for key in target_dicts[0].keys():
138+ tensors = [d[key] for d in target_dicts if key in d and d[key] is not None]
139+ if not tensors:
140+ continue
141+ target_tensors[key] = torch.stack(tensors)
142+ return input_tensors, target_tensors
143+ 
144+ 
145+class HDF5Dataset(Dataset):
146+ def __init__(self, hdf5_path):
147+ self.hdf5_path = hdf5_path
148+ self._load_hdf5()
149+ 
150+ def _load_hdf5(self):
151+ with h5py.File(self.hdf5_path, 'r') as f:
152+ self.input_sample = {}
153+ self.target_sample = {}
154+ y = np.array(f["y"])
155+ z = np.array(f["z"])
156+ fields = [key for key in f.keys() if key not in ["y", "z"]]
157+ self.target_sample.update({"y": torch.tensor(y, dtype=torch.float32)})
158+ self.target_sample.update({"z": torch.tensor(z, dtype=torch.float32)})
159+ for multi_field in fields:
160+ self.input_sample.update({multi_field: torch.tensor(np.array(f[multi_field]), dtype=torch.int64)})
161+ self._length = len(y)
162+ self.positive_indices = np.where(y == 1)[0]
163+ self.negative_indices = np.where(y == 0)[0]
164+ # logger.info(f"load file {self.hdf5_path} finished")
165+ 
166+ def __len__(self):
167+ return self._length
168+ 
169+ def __getitem__(self, idx):
170+ input_dict = {k: v[idx] for k, v in self.input_sample.items()}
171+ target_dict = {k: v[idx] for k, v in self.target_sample.items()}
172+ return input_dict, target_dict
173+ 
174+ 
175+class TorchDataSet(ConcatDataset):
176+ def __init__(self, files):
177+ print(files)
178+ datasets = [HDF5Dataset(fp) for fp in files]
179+ super().__init__(datasets)
180+ 
181+ 
182+class TorchMmoeModel(nn.Module):
183+ def __init__(self, params) -> None:
184+ super().__init__()
185+ self.params = params
186+ self.spec = json_file_load("spec", os.path.join(self.params.data_dir, "spec.json"))
187+ self.embedding_layers = self.build_embedding_layers()
188+ self.experts = self.build_experts()
189+ self.gates = self.build_gate_networks()
190+ self.task_output_layers = self.build_task_output_layers()
191+ 
192+ self.tower_names = ['ctr', 'cvr']
193+ self.towers = nn.ModuleDict()
194+ self.towers_output_layers = nn.ModuleDict()
195+ tower_units = list(map(int, self.params.tower_layers.strip().split(',')))
196+ input_dim = self.params.experts_num * list(map(int, self.params.expert_layers.strip().split(',')))[-1]
197+ for name in self.tower_names:
198+ tower_layers = []
199+ in_dim = input_dim
200+ for out_dim in tower_units:
201+ tower_layers.append(nn.Linear(in_dim, out_dim))
202+ tower_layers.append(nn.BatchNorm1d(out_dim))
203+ tower_layers.append(nn.ReLU())
204+ tower_layers.append(nn.Dropout(p=0.2))
205+ in_dim = out_dim
206+ self.towers[name] = nn.Sequential(*tower_layers)
207+ self.towers_output_layers[name] = nn.Linear(in_dim, 1)
208+ 
209+ def forward(self, features: dict):
210+ # Build the embedding layer
211+ x_deep = self.get_embedding(features)
212+ experts_out = [expert(x_deep) for expert in self.experts]
213+ experts_out = torch.stack(experts_out, dim=1)
214+ 
215+ gate_outs = [gate(x_deep) for gate in self.gates]
216+ 
217+ task_outputs = []
218+ for gate_network in gate_outs:
219+ gate_network = gate_network.unsqueeze(-1)
220+ task_out = torch.multiply(experts_out, gate_network)
221+ task_out_shape = list(task_out.shape)
222+ task_outputs.append(torch.reshape(task_out, shape=[-1, task_out_shape[1] * task_out_shape[2]]))
223+ return task_outputs
224+ 
225+ def build_embedding_layers(self):
226+ embeddings = nn.ModuleDict()
227+ for key, vocab_len in self.spec["vocab_length"].items():
228+ embeddings[key] = nn.Embedding(vocab_len + 1, self.params.embedding_size)
229+ nn.init.normal_(embeddings[key].weight, std=STD_DEV)
230+ return embeddings
231+ 
232+ def build_experts(self):
233+ 
234+ experts = nn.ModuleList()
235+ expert_units = list(map(int, self.params.expert_layers.strip().split(',')))
236+ input_dim = self.params.embedding_size * EMBEDDING_FEATURE_NUM
237+ 
238+ for _ in range(self.params.experts_num):
239+ expert_layers = []
240+ in_features = input_dim
241+ 
242+ for out_features in expert_units:
243+ expert_layers.append(nn.Linear(in_features, out_features))
244+ expert_layers.append(nn.BatchNorm1d(out_features))
245+ expert_layers.append(nn.ReLU())
246+ expert_layers.append(nn.Dropout(0.2))
247+ in_features = out_features
248+ experts.append(nn.Sequential(*expert_layers))
249+ return experts
250+ 
251+ def build_gate_networks(self):
252+ gates = nn.ModuleList()
253+ input_dim = self.params.embedding_size * EMBEDDING_FEATURE_NUM
254+ 
255+ for _ in range(self.params.task_num):
256+ gate = nn.Sequential(
257+ nn.Linear(input_dim, self.params.experts_num),
258+ nn.Softmax(dim=1)
259+ )
260+ gates.append(gate)
261+ return gates
262+ 
263+ def build_task_output_layers(self):
264+ task_output_layers = nn.ModuleList()
265+ expert_units = list(map(int, self.params.expert_layers.strip().split(',')))
266+ input_dim = self.params.experts_num * expert_units[-1]
267+ tower_units = list(map(int, self.params.tower_layers.strip().split(',')))
268+ for _ in range(self.params.task_num):
269+ tower = []
270+ in_features = input_dim
271+ for out_features in tower_units:
272+ tower.append(nn.Linear(in_features, out_features))
273+ tower.append(nn.ReLU())
274+ in_features = out_features
275+ task_output_layers.append(nn.Sequential(*tower))
276+ return task_output_layers
277+ 
278+ def embedding_lookup_sparse_fake(self, key,
279+ ids: torch.Tensor,
280+ combiner: str = None,
281+ name: str = None) -> torch.Tensor:
282+ dense_mask = torch.unsqueeze(torch.where(ids >= 0,
283+ torch.ones_like(ids, dtype=torch.float32),
284+ torch.zeros_like(ids)
285+ ),
286+ dim=-1
287+ )
288+ 
289+ # Replace invalid IDs (-1) with zeros
290+ ids = torch.where(ids == -1, torch.zeros_like(ids), ids)
291+ embedding_layer = self.embedding_layers[key]
292+ embedding_output = embedding_layer(ids)
293+ embedding = embedding_output * dense_mask
294+ summed_embedding = torch.sum(embedding, axis=1)
295+ if combiner == "sum":
296+ return summed_embedding
297+ elif combiner == "mean":
298+ return summed_embedding / torch.sum(dense_mask, axis=1)
299+ else:
300+ raise ValueError("combiner only supoort 'sum', 'mean'")
301+ 
302+ def get_embedding(self, features: dict) -> torch.Tensor:
303+ """
304+ Build the embedding layer for the model.
305+ 
306+ Args:
307+ features (dict): The input features.
308+ Returns:
309+ torch.Tensor: The concatenated and reshaped embedding tensor.
310+ """
311+ on_hot_field_lst = self.spec.get("one_hot_fields")
312+ other_field_lst = self.spec.get("multi_hot_fields") + self.spec.get("special_fields")
313+ 
314+ embeddings = {}
315+ for key in on_hot_field_lst:
316+ embedding_layer = self.embedding_layers[key]
317+ embeddings[key] = embedding_layer(features[key])
318+ embeddings[key] = torch.reshape(embeddings[key], [-1, 1, self.params.embedding_size])
319+ for key in other_field_lst:
320+ embeddings[key] = torch.unsqueeze(
321+ self.embedding_lookup_sparse_fake(key=key, ids=features[key], combiner="sum"),
322+ dim=1
323+ )
324+ embedding = torch.concat(
325+ [embeddings.get(field_name) for field_name in self.spec.get("one_hot_fields")] +
326+ [embeddings.get(field_name) for field_name in self.spec.get("multi_hot_fields")] +
327+ [embeddings.get(field_name) for field_name in self.spec.get("special_fields")],
328+ dim=2,
329+ )
330+ 
331+ return torch.reshape(embedding, [-1, EMBEDDING_FEATURE_NUM * self.params.embedding_size])
332+ 
333+ def build_tower(self, tower_input: torch.Tensor, name: str) -> torch.Tensor:
334+ """
335+ Build the tower network for a specific task.
336+ 
337+ Args:
338+ tower_input (torch.Tensor): The input tensor for the tower.
339+ name (str): The name of the tower.
340+ Returns:
341+ torch.Tensor: The output tensor of the tower network.
342+ """
343+ tower_units = list(map(int, self.params.tower_layers.strip().split(',')))
344+ y_tower = tower_input
345+ for tower_i, _ in enumerate(tower_units):
346+ tower_linear = nn.Linear(in_features=y_tower.shape[-1], out_features=tower_units[tower_i])
347+ tower_linear_out = tower_linear(y_tower)
348+ relu = nn.ReLU()
349+ y_tower = relu(tower_linear_out)
350+ return y_tower
351+ 
352+ def build_predictions(self, task_outputs: list) -> dict:
353+ """
354+ Build the predictions for the model.
355+ 
356+ Args:
357+ task_outputs (list): A list of task output tensors.
358+ 
359+ Returns:
360+ dict: A dictionary containing the predictions for ctr, cvr, and ctcvr.
361+ """
362+ preds = {}
363+ tower_outputs = {}
364+ for i, name in enumerate(self.tower_names):
365+ y = self.towers[name](task_outputs[i])
366+ y = self.towers_output_layers[name](y)
367+ y = torch.reshape(y, [-1, ])
368+ preds[name + '_logit'] = y
369+ preds[name] = torch.sigmoid(y)
370+ tower_outputs[name] = y
371+ 
372+ ctr_pred = preds.get('ctr')
373+ cvr_pred = preds.get('cvr')
374+ if ctr_pred is not None and cvr_pred is not None:
375+ if 'ctr_logit' not in preds or 'cvr_logit' not in preds:
376+ raise ValueError("Missing required keys in preds dictionary")
377+ else:
378+ preds['ctcvr'] = ctr_pred * cvr_pred
379+ preds['ctcvr_logit'] = preds['ctr_logit'] + preds['cvr_logit']
380+ return preds
381+ 
382+ def build_loss(self,
383+ labels: dict,
384+ y_ctr_logit: torch.Tensor,
385+ y_ctcvr_logit: torch.Tensor,
386+ ctr_weight,
387+ ctcvr_weight) -> torch.Tensor:
388+ """
389+ Build the loss function for the model.
390+ 
391+ Args:
392+ labels (dict): A dictionary containing the true labels for ctr and ctcvr.
393+ y_ctr_logit (torch.Tensor): The predicted ctr values.
394+ y_ctcvr_logit (torch.Tensor): The predicted ctcvr values.
395+ Returns:
396+ torch.Tensor: The combined loss tensor.
397+ """
398+ y_ctr_logit = y_ctr_logit.view(-1)
399+ y_ctcvr_logit = y_ctcvr_logit.view(-1)
400+ y = labels['y'].view(-1)
401+ z = labels['z'].view(-1)
402+ 
403+ bce_loss = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([ctr_weight],
404+ dtype=torch.float32,
405+ device=y.device)
406+ )
407+ ctcvr_bce_loss = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([ctcvr_weight],
408+ dtype=torch.float32,
409+ device=z.device)
410+ )
411+ 
412+ ctr_loss = bce_loss(y_ctr_logit, y)
413+ ctcvr_loss = ctcvr_bce_loss(y_ctcvr_logit, z)
414+ ctr_task_wgt = self.params.ctr_task_wgt
415+ 
416+ return ctr_task_wgt * ctr_loss + (1 - ctr_task_wgt) * ctcvr_loss
417+ 
418+ def build_optimizer(self):
419+ """
420+ Build the optimizer for training.
421+ 
422+ Args:
423+ loss (torch.Tensor): The loss tensor to minimize.
424+ 
425+ Returns:
426+ optim: The operation for applying gradients.
427+ 
428+ Raises:
429+ ValueError: If the optimizer type is not supported.
430+ """
431+ if self.params.optimizer == "Adam":
432+ optimizer = optim.Adam(
433+ params=self.parameters(),
434+ lr=self.params.learning_rate,
435+ betas=[0.9, 0.99], eps=1e-8
436+ )
437+ elif self.params.optimizer == "Adagrad":
438+ optimizer = optim.Adagrad(
439+ params=self.parameters(),
440+ lr=self.params.learning_rate,
441+ initial_accumulator_value=1e-6
442+ )
443+ elif self.params.optimizer == "Momentum":
444+ optimizer = optim.SGD(
445+ params=self.parameters(),
446+ lr=self.params.learning_rate,
447+ momentum=0.95
448+ )
449+ elif self.params.optimizer == "SGD":
450+ optimizer = optim.SGD(
451+ params=self.parameters(),
452+ lr=self.params.learning_rate,
453+ )
454+ else:
455+ raise ValueError("Unsupported optimizer type: {}".format(args.optimizer))
456+ return optimizer
457+ 
458+ 
459+def get_profile(profiler_start_step: int, profiler_end_step: int, profiling_save_path: str):
460+ warm_step = profiler_start_step
461+ active_step = profiler_end_step - warm_step +1
462+ print(f"[Profile INFO]: warm_step: {warm_step}, active_step: {active_step}, profiling_save_path: {profiling_save_path}")
463+ device = detect_device_type()
464+ if device == 'npu':
465+ import torch_npu
466+ g_prof_config = torch_npu.profiler._ExperimentalConfig(
467+ export_type=[
468+ torch_npu.profiler.ExportType.Text,
469+ torch_npu.profiler.ExportType.Db
470+ ],
471+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
472+ msprof_tx=False,
473+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
474+ l2_cache=False,
475+ op_attr=False,
476+ data_simplification=False,
477+ record_op_args=False,
478+ gc_detect_threshold=None)
479+ 
480+ return torch_npu.profiler.profile(
481+ activities=[
482+ torch_npu.profiler.ProfilerActivity.CPU,
483+ torch_npu.profiler.ProfilerActivity.NPU],
484+ schedule=torch_npu.profiler.schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
485+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profiling_save_path),
486+ record_shapes=False,
487+ profile_memory=False,
488+ with_stack=False,
489+ with_modules=False,
490+ with_flops=False,
491+ experimental_config=g_prof_config)
492+ elif device == 'cuda':
493+ from torch.profiler import profile, schedule, tensorboard_trace_handler, ProfilerActivity
494+ return profile(
495+ activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
496+ schedule=schedule(wait=0, warmup=warm_step, active=active_step, repeat=1),
497+ on_trace_ready=tensorboard_trace_handler(profiling_save_path),
498+ record_shapes=False,
499+ profile_memory=False,
500+ )
501+ else:
502+ print(f"[Warning]: No supported acceleration device (CUDA/NPU) detected, profiler will be disabled")
503+ return None
504+ 
505+ 
506+def device_synchronize():
507+ if torch.cuda.is_available():
508+ torch.cuda.synchronize()
509+ elif torch.npu.is_available():
510+ torch.npu.synchronize()
511+ 
512+ 
513+def evaluate_op(model: TorchMmoeModel, te_files, device, args):
514+ te_dataset = TorchDataSet(te_files)
515+ te_ctr, te_ctcvr = pre_deal_dataset(te_dataset)
516+ test_dataloader = DataLoader(dataset=te_dataset,
517+ batch_size=args.batch_size,
518+ shuffle=True,
519+ collate_fn=collate_fn,
520+ prefetch_factor=100,
521+ num_workers=10)
522+ dl = iter(test_dataloader)
523+ input_sample, target_sample = next(dl)
524+ input_sample = {k: v.to(device) for k, v in input_sample.items()}
525+ 
526+ mod = 'compile' if args.enable_compile else 'eager'
527+ if mod == 'compile':
528+ model = torch.compile(model, mode='reduce-overhead', dynamic=False)
529+ 
530+ model.eval()
531+ with torch.no_grad():
532+ model(input_sample)
533+ device_synchronize()
534+ 
535+ prof = None
536+ if args.enable_profiler:
537+ profiling_save_path = args.profiler_save_path + '/' + MODEL_NAME + '/' + mod
538+ prof = get_profile(args.profiler_start_step, args.profiler_end_step, profiling_save_path)
539+ prof.start()
540+ 
541+ exec_times = []
542+ with torch.no_grad():
543+ for i in range(args.max_steps):
544+ e2etime1 = time.perf_counter()
545+ model(input_sample)
546+ device_synchronize()
547+ e2etime2 = time.perf_counter()
548+ elapsed_ms = (e2etime2 - e2etime1) * 1000
549+ print(f"iterations {i} / {args.max_steps}: [{mod}] e2e time: {elapsed_ms} ms.")
550+ 
551+ if i >= 10:
552+ exec_times.append(elapsed_ms)
553+ 
554+ if args.enable_profiler:
555+ prof.step()
556+ if args.enable_profiler:
557+ prof.stop()
558+ if exec_times:
559+ avg_time = sum(exec_times) / len(exec_times)
560+ print("Step time consumption statistics (excluding the first 10 steps)")
561+ print(f"\n>>> [{mod}] avg e2e: {avg_time:.3f} ms <<<")
562+ 
563+ 
564+def main(args):
565+ model = TorchMmoeModel(args)
566+ te_files = glob.glob("%stest/data_test.csv.hd5.*" % args.data_dir)
567+ device_type = detect_device_type()
568+ print(f"Using device type {device_type}")
569+ device = torch.device(device_type)
570+ model.to(device)
571+ evaluate_op(model, te_files, device, args)
572+ 
573+ 
574+if __name__ == "__main__":
575+ print(f"{MODEL_NAME} inference begin !!!")
576+ args = parse_arguments()
577+ main(args)
Abenchmarks/inference/README.md+92-0
@@ -0,0 +1,92 @@
1+# 推理说明
2+
3+本 README 说明如何使用 **DCNv2**，**DIN**，**MMOE**和**ETA**模型进行推理（含 eager / torch.compile 两种模式）。
4+在运行完成后输出
5+---
6+ 
7+## 目录
8+ 
9+- [1. 数据集处理](#1-数据集处理)
10+ - [1.1 DIN & DCNv2](#11-DIN-&-DCNv2)
11+ - [1.2 MMOE & ETA](#12-MMOE-&-ETA)
12+- [2. 模型推理](#2-模型推理)
13+ - [2.1 eager mode（默认）](#21-eager-mode默认)
14+ - [2.2 启用 torch.compile（可选）](#22-启用-torch.compile可选)
15+ - [2.3 采集 profile文件（可选）](#23-采集profile文件可选)
16+---
17+ 
18+## 1. 数据集处理
19+ 
20+### 1.1 DIN & DCNv2
21+ 
22+**DIN** 和 **DCNv2** 的推理数据集来自 CTR_Algorithm 仓库示例数据，该数据集无需手动下载
23+ 
24+https://github.com/Prayforhanluo/CTR_Algorithm/tree/main/data
25+ 
26+在DIN或DCNv2目录下执行如下命令：
27+```bash
28+git clone https://github.com/Prayforhanluo/CTR_Algorithm.git
29+```
30+- 数据集目录：`CTR_Algorithm/data/data.csv`
31+ 
32+ 
33+### 1.2 MMOE & ETA
34+ 
35+进入**MMOE**或**ETA**文件夹下，进行如下数据预处理
36+ 
37+推理数据集来自 https://tianchi.aliyun.com/dataset/408 Ali-CPP数据集，需做如下处理：
38+1. 下载数据集，存放至当前目录下alicpp文件夹中
39+2. 后续流程参考：https://gitee.com/ascend/RecSDK/blob/develop_torch_benchmark/torch_examples_benchmark/model_zoo/README.md
40+ 
41+执行完成后数据集会默认生成至aliccp_out目录下。
42+ 
43+网络依赖已整理到 `requirements.txt`，可直接安装：
44+ 
45+```bash
46+pip install -r ./requirements.txt
47+```
48+ 
49+在模型执行时，默认加载数据集路径为: `./aliccp/aliccp_out/`；模型脚本同时支持运行时使用如下命令动态指定数据集目录：
50+ 
51+```bash
52+python eta.py --data_dir path/to/your/data/
53+```
54+ 
55+ 
56+## 2. 模型推理(DCNv2网络为例)
57+ 
58+进入**DCNv2**文件夹，直接执行推理脚本`dcnv2.py`即可
59+ 
60+- 注：运行该网络前需要先将patch文件加上，具体方式如下：
61+ ```bash
62+ cd DCNv2
63+ unix2dos dcnv2.patch #如果在arm机器上需要执行
64+ git apply dcnv2.patch
65+ ```
66+ 
67+### 2.1 eager mode（默认）
68+ 
69+```bash
70+python dcnv2.py
71+```
72+ 
73+### 2.2 启用 `torch.compile`（可选）
74+ 
75+脚本已支持 `--enable_compile` 这样的开关，开启方式为：
76+ 
77+```bash
78+python dcnv2.py \
79+--enable_compile
80+```
81+ 
82+### 2.3 采集profile文件（可选）
83+ 
84+脚本已支持 `--enable_profiler` 这样的开关，开启方式为：
85+ 
86+```bash
87+python dcnv2.py \
88+ --enable_profiler \
89+ --profiler_start_step 5 \
90+ --profiler_end_step 6 \
91+```
92+可以通过 `--profiler_start_step ` 和 `--profiler_end_step` 分别设置profile开始和结束步数。
Abenchmarks/inference/requirements.txt+5-0
@@ -0,0 +1,5 @@
1+h5py==3.13.0
2+numpy==1.26.4
3+pandas
4+pytz
5+scikit-learn