├── GA_exp.py
├── LICENSE
├── README.md
├── arguments.py
├── critic_train_exp.py
├── layouts.py
├── multi_scales_test.py
├── off_loading_models.py
└── pretrain.py


/GA_exp.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import random
  3 | from tqdm import tqdm
  4 | from arguments import args
  5 | import numpy as np
  6 | from layouts import generate_layouts
  7 | from torch_geometric.loader import DataLoader
  8 | from torch_geometric.utils import softmax
  9 | from time import time
 10 | 
 11 | 
 12 | 
 13 | class Evolution():
 14 |     def __init__(self, N, M, pop_num, total_epochs, compute_resource, path_losses, task_size, edge_index) -> None:
 15 |         
 16 |         self.N = N      # user数量
 17 |         self.M = M      # server数量
 18 |         self.pop_num = pop_num     # 初始种群规模
 19 |         self.retain_rate = 0.4      # 保存率
 20 |         self.mutate_rate = 0.2
 21 |         self.random_select_rate = 0.1
 22 |         # self.locations = locations
 23 |         self.total_epoch = total_epochs
 24 |         self.b = 2
 25 | 
 26 | 
 27 |         self.compute_resource = compute_resource
 28 |         self.path_losses = path_losses
 29 |         self.task_size = task_size
 30 |         self.edge_index = edge_index
 31 | 
 32 |     def evolution(self):
 33 |         populication = self.populication_create()
 34 |         loss_list = []
 35 |         time_list = []
 36 |         for i in range(self.total_epoch):
 37 |             # print(populication.shape)
 38 |             start = time()
 39 |             parents, output_i = self.selection(populication)
 40 |             cs = self.cross_over(parents)
 41 |             cs = self.mutation(cs, i)
 42 |             populication = torch.cat([parents, cs], dim=0)
 43 |             # output_i = self.adaptbility(populication)
 44 |             
 45 |             min_time_loss = torch.min(output_i)       # 最好的一个对应的time loss
 46 |             
 47 |             loss_list.append(min_time_loss)
 48 |             # print('epoch == {}, time of best_one == {}'.format(i, min_time_loss))
 49 |             end = time()
 50 |             time_list.append(end-start)
 51 |         loss_list = torch.stack(loss_list)
 52 |         time_list = np.array(time_list)
 53 |         return loss_list, time_list
 54 |     
 55 |         
 56 | 
 57 |     def cross_over(self, parent):
 58 |         # 交叉, 单点交叉
 59 |         '''
 60 |         #选择父代
 61 |         male = []
 62 |         female = []
 63 |         # 选择基因位并交叉      //单点交叉
 64 |         '''
 65 | 
 66 |         # 均匀交叉
 67 |         children = []
 68 |         get_child_num = self.pop_num-len(parent)
 69 |         while len(children) < get_child_num:
 70 |             i = random.randint(0, len(parent)-1)
 71 |             j = random.randint(0, len(parent)-1)
 72 |             male = parent[i]
 73 |             female = parent[j]
 74 |             select_p = torch.rand(len(male), device=args.device)
 75 |             select_p[torch.where(select_p < 0.5)] = 0
 76 |             select_p[torch.where(select_p >= 0.5)] = 1
 77 |             child1 = select_p * male + (1-select_p) * female
 78 |             child2 = (1 - select_p) * male + select_p * female
 79 |             children.append(child1.reshape(1, len(child1)))
 80 |             children.append(child2.reshape(1, len(child2)))
 81 |         if len(children) != 0:
 82 |             children = torch.cat(children, dim=0)
 83 |         if get_child_num < len(children):
 84 |             children = children[:-1]
 85 |         return children
 86 | 
 87 |     def populication_create(self):
 88 |         # 生成种群
 89 |         self.populication = torch.rand((self.pop_num, 3*self.M*self.N), device=args.device)
 90 |         # self.users = torch.tensor(self.features[:2*self.N], device=args.device)
 91 |         
 92 |         return self.populication
 93 | 
 94 |     def mutation(self, cs, i):
 95 |         # 变异
 96 |         
 97 |         # 采用非一致性变异，每个位置都进行变异
 98 |         new_cs = cs.clone()
 99 |         for idx, c in enumerate(cs):
100 |             if random.random() < self.mutate_rate:
101 |                 r = random.random()
102 |                 mut1 = (1-c)*torch.rand(len(c), device=args.device)*(1-i/self.total_epoch)**self.b
103 |                 mut2 = torch.rand(len(c), device=args.device)*(1-i/self.total_epoch)**self.b
104 |                 # print(mut1)
105 |                 if random.random() > 0.5:
106 |                     c = c + mut1
107 |                 else:
108 |                     c = c - mut2
109 |                 # print(c)
110 |             new_cs[idx] = c
111 |             # print(c)
112 |         return new_cs
113 |             
114 | 
115 |     def selection(self, populication):
116 |         # 选择
117 | 
118 |         # 选择最佳的rate率的个体
119 |         # 对种群从小到大进行排序
120 |         adpt = self.adaptbility(populication)
121 |         # grabed = [[ad, one] for ad, one in zip(adpt, populication)]
122 |         
123 |         sort_index = torch.argsort(adpt)
124 |         grabed = populication[sort_index]
125 |         sorted_adpt = adpt[sort_index]
126 |         # sorted_grabed = sorted(grabed, key=lambda x: x[0])
127 |         # grabed = torch.tensor([x[1] for x in sorted_grabed], device=args.device)
128 |         index = int(len(populication)*self.retain_rate)
129 | 
130 |         live = grabed[:index]
131 |         
132 |         live_adpt = sorted_adpt[:index]
133 | 
134 |         # 选择幸运个体
135 |         for i, ad_i in zip(grabed[index:], adpt[index:]):
136 |             if random.random() < self.random_select_rate:
137 |                 live = torch.cat([live, i.reshape(1, len(i))], dim=0)
138 |                 live_adpt = torch.cat([live_adpt, ad_i.unsqueeze(0)], dim=0)
139 |                 # live_adpt = torch.stack([live_adpt, ad_i.unsqueeze(0)])
140 |         
141 |         return live, adpt
142 |     
143 |     def adaptbility(self, populication):
144 | 
145 |         task_allocation = populication[:, :self.M*self.N]
146 |         task_allocation = softmax(task_allocation, index=self.edge_index[0], dim=1)
147 |         power_allocation = populication[:, self.M*self.N:2*self.M*self.N]
148 |         power_allocation = softmax(power_allocation, index=self.edge_index[0], dim=1)
149 |         comp_allocation = populication[:, 2*self.N*self.M:]
150 |         comp_allocation = softmax(comp_allocation, index=self.edge_index[1], dim=1)
151 | 
152 | 
153 |         
154 |         time_losses = self.compute_loss(task_allocation, power_allocation, comp_allocation)
155 |         
156 |         # max_dist = []
157 |         # for p in torch.FloatTensor(populication).to(args.device):
158 |         #     p = p.reshape(int(len(p)/2), 2)
159 |         #     # p = torch.FloatTensor(p).to(device)
160 |         #     users_uav = torch.cat([self.users, p], dim=0)
161 |         #     max_dist.append(self.flow_loss(users_uav).cpu().data.numpy())
162 |         return time_losses.mean(-1)
163 | 
164 |     def compute_loss(self, task_allocation, power_allocation, comp_allocation):
165 |         
166 |         # task_size : vector N
167 |         # task_allocation: mat pop_num x 3*M*N
168 |         # index: vector 3*M*N
169 |         
170 |         epsilon = 1e-9
171 |         extre = 1e-20
172 |         user_index = self.edge_index[0]      # s2u中源节点的索引
173 |         server_index = self.edge_index[1]    # s2u中目标节点的索引
174 |         
175 |         task_size = self.task_size[user_index]       # M*N    
176 |         # task_size = task_size[user_index]*args.tasksize_cof       # 重复采样映射到边中
177 |         
178 |         tasks = task_size * task_allocation   # mat pop_num x M*N
179 | 
180 |         compute_resource = self.compute_resource[server_index]
181 |         # compute_resource = compute_resource[server_index]*args.comp_cof       # 
182 | 
183 |         comp = compute_resource * comp_allocation
184 | 
185 |         pw = power_allocation * self.path_losses    # mat pop_num x M*N
186 |         # pw = torch.clamp(pw, 1e-5, 1)
187 | 
188 |         pw_list = torch.zeros((pw.shape[0], pw.shape[1], server_index[-1]+1), device=args.device)   # mat pop_num x MN x N
189 |         pw_list.scatter_(2, server_index.repeat((self.pop_num, 1)).unsqueeze(2), pw.unsqueeze(2))
190 |         pws_list = pw_list.sum(1)[:, server_index]  # mat pop_num x MN
191 | 
192 | 
193 |         interference = pws_list-pw
194 |         rate = torch.log2(1+torch.div(pw, interference+epsilon))
195 |         # rate = args.band_width * torch.log2(1+torch.div(pw, interference+epsilon))
196 |         # offloading_time = torch.div(tasks, rate+extre) * (args.tasksize_cof/args.band_width)
197 |         offloading_time = torch.div(tasks, rate+extre)
198 | 
199 |         # compute_time = torch.div(tasks, comp+extre) * (args.tasksize_cof*args.cons_factor/args.comp_cof)
200 |         compute_time = torch.div(tasks, comp+extre)
201 | 
202 |         time_loss = offloading_time + compute_time      # pop_num x MN
203 |         assert torch.isnan(time_loss).sum()==0
204 | 
205 | 
206 |         time_loss_list = torch.zeros((time_loss.shape[0], time_loss.shape[1], user_index[-1]+1), device=args.device)
207 |         time_loss_list.scatter_(2, user_index.repeat((self.pop_num, 1)).unsqueeze(2), time_loss.unsqueeze(2))
208 |         time_loss_list = time_loss_list.sum(1)      # pop_num x MN
209 | 
210 |         return time_loss_list
211 | 
212 | 
213 | if __name__=='__main__':
214 |     server_num = 5
215 | 
216 |     pop_num = 200
217 |     epochs = 600
218 | 
219 |     sample_num = 5
220 | 
221 | 
222 |     train_user_nums = np.random.randint(server_num*3, server_num*3+1, sample_num)
223 |     train_server_nums = np.random.randint(server_num, server_num+1, sample_num)
224 |     # test_user_nums = np.random.randint(server_num*3, server_num*3+1, args.test_layouts)
225 |     # test_server_nums = np.random.randint(server_num, server_num+1, args.test_layouts)
226 | 
227 |     # env_max_length = np.sqrt(server_num * 300)
228 | 
229 |     train_layouts = generate_layouts(train_user_nums, train_server_nums, args)
230 |     # test_layouts = generate_layouts(test_user_nums, test_server_nums, args)
231 | 
232 |     train_loader = DataLoader(train_layouts, batch_size=args.batch_size, shuffle=True)
233 |     # test_loader = DataLoader(test_layouts, batch_size=args.batch_size, shuffle=True)
234 |     loss_list = []
235 |     for data in train_layouts:
236 |         compute_resource = data['server'].x[:, 0].squeeze()
237 |         path_losses = data['user', 'u2s', 'server'].path_loss.squeeze()
238 |         task_size = data['user'].x[:, 0].squeeze()
239 |         edge_index = data['user', 'u2s', 'server'].edge_index
240 |         ga = Evolution(server_num, server_num*3, pop_num, epochs, compute_resource=compute_resource, path_losses=path_losses, task_size=task_size, edge_index=edge_index)
241 |         loss, _ = ga.evolution()
242 |         loss_list.append(loss.item()[-1])
243 |     print(np.mean(loss_list))
244 | 
245 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2023 UNIC Lab
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # Scalable Resource Management for Dynamic MEC: An Unsupervised Link-Output Graph Neural Network Approach
2 | This is the code for paper "Scalable Resource Management for Dynamic MEC: An Unsupervised Link-Output Graph Neural Network Approach" 
3 | [Paper](https://arxiv.org/pdf/2306.08938.pdf)
4 | 


--------------------------------------------------------------------------------
/arguments.py:
--------------------------------------------------------------------------------
 1 | 
 2 | import argparse
 3 | args = argparse.ArgumentParser()
 4 | #offloading 模型参数
 5 | args.add_argument('--learning_rate', default=1e-4)
 6 | args.add_argument('--critic_lr', default=5e-4)
 7 | args.add_argument('--input_dim', default=2)
 8 | args.add_argument('--state_dim', default=3)
 9 | args.add_argument('--action_dim', default=2)
10 | args.add_argument('--hidden_dim', default=64)
11 | args.add_argument('--alpha', default=0.2)
12 | args.add_argument('--device', default='cuda:0')
13 | args.add_argument('--load_pretrained', default=False)
14 | args.add_argument('--num_layers', default=2)
15 | 
16 | # 实验环境参数
17 | args.add_argument('--user_num', default=20)
18 | args.add_argument('--server_num', default=5)
19 | args.add_argument('--test_user_num', default=20)
20 | args.add_argument('--test_server_num', default=5)
21 | args.add_argument('--p_max', default=1)
22 | args.add_argument('--pw_threshold', default=1e-6)
23 | args.add_argument('--train_layouts', default=128)
24 | args.add_argument('--test_layouts', default=64)
25 | args.add_argument('--env_max_length', default=400)
26 | args.add_argument('--server_height', default=20)
27 | args.add_argument('--carrier_f_start', default=2.4e9)
28 | args.add_argument('--carrier_f_end', default=2.4835e9)
29 | args.add_argument('--signal_cof', default=4.11)
30 | args.add_argument('--band_width', default=1e6)
31 | args.add_argument('--batch_size', default=32)
32 | args.add_argument('--max_server_num', default=15)
33 | args.add_argument('--init_min_size', default=2)
34 | args.add_argument('--init_max_size', default=8)
35 | 
36 | args.add_argument('--cons_factor', default=10)
37 | args.add_argument('--init_min_comp', default=0.1)
38 | args.add_argument('--init_max_comp', default=1)
39 | args.add_argument('--comp_cof', default=1024**2)
40 | args.add_argument('--tasksize_cof', default=1024*100)
41 | 
42 | 
43 | 
44 | args.add_argument('--multi_scales_train', default=False)
45 | 
46 | args.add_argument('--multi_scales_test', default=False)
47 | 
48 | args.add_argument('--single_scale_test', default=True)
49 | 
50 | args.add_argument('--comparison_hgnn', default=True)
51 | args.add_argument('--comparison_pcnet', default=True)
52 | args.add_argument('--comparison_pcnetCritic', default=False)
53 | 
54 | 
55 | args.add_argument('--train_steps', default=600)
56 | args.add_argument('--evaluate_steps', default=10)
57 | args.add_argument('--save_steps', default=50)
58 | 
59 | 
60 | 
61 | args = args.parse_args()
62 | 
63 | 


--------------------------------------------------------------------------------
/critic_train_exp.py:
--------------------------------------------------------------------------------
  1 | 
  2 | import numpy as np
  3 | import torch
  4 | import torch.nn as nn
  5 | import matplotlib.pyplot as plt
  6 | from off_loading_models import PCNet, PCNetCritic, MMSE, GnnCritic, TaskLoad
  7 | from arguments import args
  8 | from torch_geometric.loader import DataLoader
  9 | from layouts import generate_layouts
 10 | from tqdm import tqdm
 11 | 
 12 | 
 13 | def compute_loss_nn(task_allocation, power_allocation, comp_allocation, task_size, compute_resource, path_losses, user_index, server_index):
 14 |     
 15 |     # task_size : vector N
 16 |     # task_allocation: mat pop_num x 3*M*N
 17 |     # index: vector 3*M*N
 18 |     
 19 |     epsilon = 1e-9
 20 |     extre = 1e-20
 21 |     server_index_first = server_index.reshape((batch_size, -1))[0]
 22 |     user_index_first = user_index.reshape((batch_size, -1))[0]
 23 |     # user_index = edge_index[0]      # s2u中源节点的索引
 24 |     # server_index = edge_index[1]    # s2u中目标节点的索引
 25 |     
 26 |     # power_allocation = torch.clamp(power_allocation, 1e-5, 1)
 27 |     pw_ini = power_allocation * path_losses    # mat pop_num x M*N
 28 |     
 29 |     # 将信道状态过小的设置为0
 30 |     mask_pw = torch.where(pw_ini<args.pw_threshold)
 31 |     
 32 |     pw = pw_ini.clone()
 33 | 
 34 |     pw[mask_pw] = 0
 35 | 
 36 |     comp_allocation_clone = comp_allocation.clone()
 37 |     comp_allocation_clone[mask_pw] = 0
 38 |     comp_allocation_normed = torch.zeros((comp_allocation_clone.shape[0], comp_allocation_clone.shape[1], server_index_first[-1]+1), device=args.device)
 39 |     comp_allocation_normed.scatter_(2, server_index_first.repeat((batch_size, 1)).unsqueeze(2), comp_allocation_clone.unsqueeze(2))
 40 |     comp_allocation_normed = comp_allocation_normed.sum(1)[:, server_index_first]
 41 |     comp_allocation_normed = torch.div(comp_allocation_clone, comp_allocation_normed+extre)
 42 | 
 43 |     task_allocation_clone = task_allocation.clone()
 44 |     task_allocation_clone[mask_pw] = 0
 45 |     task_allocation_normed = torch.zeros((task_allocation_clone.shape[0], task_allocation_clone.shape[1], user_index_first[-1]+1), device=args.device)
 46 |     task_allocation_normed.scatter_(2, user_index_first.repeat((batch_size, 1)).unsqueeze(2), task_allocation_clone.unsqueeze(2))
 47 |     task_allocation_normed = task_allocation_normed.sum(1)[:, user_index_first]
 48 |     task_allocation_normed = torch.div(task_allocation_clone, task_allocation_normed+extre)
 49 | 
 50 |     # 计算速率
 51 |     pw_list = torch.zeros((pw.shape[0], pw.shape[1], server_index_first[-1]+1), device=args.device)   # mat pop_num x MN x N
 52 |     pw_list.scatter_(2, server_index_first.repeat((batch_size, 1)).unsqueeze(2), pw.unsqueeze(2))
 53 |     pws_list = pw_list.sum(1)[:, server_index_first]  # mat pop_num x MN
 54 | 
 55 |     interference = pws_list-pw
 56 |     rate = torch.log2(1+torch.div(pw, interference+epsilon))
 57 |     # rate = args.band_width * torch.log2(1+torch.div(pw, interference+epsilon))
 58 | 
 59 | 
 60 |     task_size = task_size[:, user_index_first]       # M*N    
 61 |     # task_size = task_size[user_index]*args.tasksize_cof       # 重复采样映射到边中
 62 |     tasks = task_size * task_allocation_normed   # mat pop_num x M*N
 63 | 
 64 |     compute_resource = compute_resource[:, server_index_first]
 65 |     # compute_resource = compute_resource[server_index]*args.comp_cof       # 
 66 | 
 67 |     comp = compute_resource * comp_allocation_normed
 68 | 
 69 | 
 70 |     # offloading_time = torch.div(tasks, rate+extre) * (args.tasksize_cof/args.band_width)
 71 |     offloading_time = torch.div(tasks, rate+extre)
 72 | 
 73 |     # compute_time = torch.div(tasks, comp+extre) * (args.tasksize_cof*args.cons_factor/args.comp_cof)
 74 |     compute_time = torch.div(tasks, comp+extre)
 75 | 
 76 |     time_loss = offloading_time + compute_time      # pop_num x MN
 77 |     assert torch.isnan(time_loss).sum()==0
 78 | 
 79 | 
 80 |     time_loss_list = torch.zeros((time_loss.shape[0], time_loss.shape[1], user_index_first[-1]+1), device=args.device)
 81 |     time_loss_list.scatter_(2, user_index_first.repeat((batch_size, 1)).unsqueeze(2), time_loss.unsqueeze(2))
 82 |     time_loss_list = time_loss_list.sum(1)      # pop_num x MN
 83 | 
 84 |     return time_loss_list.mean()
 85 | 
 86 | 
 87 | def compute_loss(task_allocation, power_allocation, comp_allocation, compute_resource, path_losses, task_size, user_index, server_index):
 88 |     
 89 |     epsilon = 1e-9
 90 |     extre = 1e-20
 91 |     # user_index = edge_index[0]      # s2u中源节点的索引
 92 |     # server_index = edge_index[1]    # s2u中目标节点的索引
 93 |     
 94 |     pw_ini = power_allocation.squeeze() * path_losses.squeeze()
 95 |     # pw小于阈值的对应power设为0
 96 |     pw = pw_ini.clone()
 97 |     # mask_pw = torch.where(pw_ini<args.pw_threshold)
 98 |     # pw[mask_pw] = 0
 99 | 
100 |     pw_user_list = torch.zeros((len(pw), user_index[-1]+1), device=args.device)
101 |     # pw_user_ini_list = torch.zeros((len(pw), user_index[-1]+1), device=args.device)
102 |     # pw_user_ini_list.scatter_(1, user_index.unsqueeze(1), pw_ini.unsqueeze(1))
103 |     pw_user_list.scatter_(1, user_index.unsqueeze(1), pw_ini.unsqueeze(1))
104 |     # 如果某一个user的发射功率均位于阈值以下
105 |     pw_masked = pw_user_list.clone()
106 |     pw_masked[torch.where(pw_masked < args.pw_threshold)] = 0
107 |     invalid_index = torch.where(pw_masked.sum(0)==0)[0]   # 是否有对所有server都低于阈值的
108 |     # assert len(invalid_index)==0
109 |     max_pw_index = pw_user_list[:, invalid_index].argmax(0) # 对所有server的pw都低于阈值的user 取信号最强的server
110 |     pw_masked[max_pw_index, invalid_index] = pw_user_list[max_pw_index, invalid_index]
111 |     pw = pw_masked.sum(1)
112 |     mask_pw = torch.where(pw==0)
113 | 
114 |     pw_list = torch.zeros((len(pw), server_index[-1]+1), device=args.device)
115 |     pw_list.scatter_(1, server_index.unsqueeze(1), pw.unsqueeze(1))
116 |     
117 |     pws_list = pw_list.sum(0)[server_index]
118 | 
119 | 
120 |     interference = pws_list-pw
121 |     rate = torch.log2(1+torch.div(pw, interference+epsilon))
122 | 
123 |     task_allocation_clone = task_allocation.clone().squeeze() + 1e-8
124 |     task_allocation_clone[mask_pw] = 0
125 |     task_allocation_normed = torch.zeros((len(task_allocation_clone), user_index[-1]+1), device=args.device)
126 |     task_allocation_normed.scatter_(1, user_index.unsqueeze(1), task_allocation_clone.unsqueeze(1))
127 |     # assert len(torch.where(task_allocation_normed.sum(0)==0)[0]) == 0
128 |     task_allocation_normed_2 = task_allocation_normed.sum(0)[user_index]
129 |     task_allocation_final = torch.div(task_allocation_clone, task_allocation_normed_2+extre)
130 |     # task_allocation_clone =  softmax(task_allocation_clone, user_index)
131 | 
132 |     task_size = task_size[user_index]
133 |     # task_size = task_size[user_index]*args.tasksize_cof       # 重复采样映射到边中
134 |     
135 |     tasks = task_size * task_allocation_final
136 | 
137 |     comp_allocation_clone = comp_allocation.clone().squeeze()
138 |     comp_allocation_clone[mask_pw] = 0
139 |     comp_allocation_normed = torch.zeros((len(comp_allocation_clone), server_index[-1]+1), device=args.device)
140 |     comp_allocation_normed.scatter_(1, server_index.unsqueeze(1), comp_allocation_clone.unsqueeze(1))
141 |     comp_allocation_normed_2 = comp_allocation_normed.sum(0)[server_index]
142 |     comp_allocation_final = torch.div(comp_allocation_clone, comp_allocation_normed_2+extre)
143 |     # compute_resource = compute_resource[server_index]*args.comp_cof       # 
144 | 
145 |     compute_resource = compute_resource[server_index]
146 |     comp = compute_resource * comp_allocation_final
147 | 
148 | 
149 |     # offloading_time = torch.div(tasks, rate+extre) * (args.tasksize_cof/args.band_width)
150 |     offloading_time = torch.div(tasks, rate+extre)
151 | 
152 |     # compute_time = torch.div(tasks, comp+extre) * (args.tasksize_cof*args.cons_factor/args.comp_cof)
153 |     compute_time = torch.div(tasks, comp+extre)
154 | 
155 | 
156 |     time_loss = offloading_time + compute_time
157 |     assert torch.isnan(time_loss).sum()==0
158 | 
159 |     time_loss_list = torch.zeros((len(time_loss), user_index[-1]+1), device=args.device)
160 |     time_loss_list.scatter_(1, user_index.unsqueeze(1), time_loss.unsqueeze(1))
161 |     time_loss_list = time_loss_list.sum(0)
162 | 
163 |     return time_loss_list.mean()
164 | 
165 | 
166 | 
167 | def gnnCritic_train(model, train_layouts, batch_size):
168 |     sum_loss = 0
169 |     length = 0
170 |     train_losses = []
171 |     test_losses = []
172 |     # actor_lr = 5e-5
173 |     # critic_lr = 1e-4
174 |     actor_optimizer = torch.optim.Adam([
175 |         {'params': model.gnn.parameters(), 'lr': actor_lr},
176 |     ])
177 |     critic_optimizer = torch.optim.Adam(model.critic_mlp.parameters(), lr=critic_lr)
178 |     critic_loss_func = torch.nn.MSELoss()
179 |     for epoch in tqdm(range(Epochs)):
180 |         model.train()
181 |         sches_list = []
182 |         for idx, batch in enumerate(train_layouts):
183 |             task_allocation, power_allocation, comp_allocation, sches, critic_value = model(batch.x_dict, batch.edge_index_dict, batch.edge_attr_dict)
184 |             user_index = batch['user', 'u2s', 'server'].edge_index[0]
185 |             server_index = batch['user', 'u2s', 'server'].edge_index[1]
186 |             total_time = compute_loss(task_allocation, power_allocation, comp_allocation, batch['server'].x[:, 0], batch['user', 'u2s', 'server'].path_loss, batch['user'].x[:, 0], user_index, server_index)
187 |             
188 |             sches_list.append(sches)
189 | 
190 |             critic_loss = critic_loss_func(total_time.detach(), critic_value)
191 |             critic_loss.backward()
192 | 
193 |             sum_loss += total_time.item()
194 |             length += 1
195 | 
196 |             if idx % batch_size == 0:
197 |                 critic_optimizer.step()
198 |                 critic_optimizer.zero_grad()
199 |                 
200 |                 sches_list = torch.vstack(sches_list)
201 |                 value_loss = model.critic_mlp(sches_list)
202 |                 sches_list = []
203 |                 value_loss.mean().backward()
204 |                 actor_optimizer.step()
205 |                 actor_optimizer.zero_grad()
206 | 
207 |                 # sum_loss = 0
208 |         
209 |         mean_loss = sum_loss/length
210 |         train_losses.append(mean_loss)
211 |         print('Epoch: {} \t\t HGnnCritic training loss: {}'.format(epoch, mean_loss))
212 |         
213 |     return train_losses
214 | 
215 | 
216 | def NN_critic_train(model, train_layouts, batch_size):
217 |     sum_loss = 0
218 |     length = 0
219 |     train_losses = []
220 |     test_losses = []
221 |     # actor_lr = 5e-5
222 |     # critic_lr = 1e-4
223 |     actor_optimizer = torch.optim.Adam([
224 |         {'params': model.encoder.parameters(), 'lr': actor_lr},
225 |         {'params': model.sche_mlp.parameters(), 'lr': actor_lr}
226 |     ])
227 |     critic_optimizer = torch.optim.Adam(model.critic_mlp.parameters(), lr=critic_lr)
228 |     critic_loss_func = torch.nn.MSELoss()
229 |     for step in tqdm(range(args.train_steps)):
230 |         model.train()
231 |         sches_list = []
232 |         for idx, batch in enumerate(train_layouts):
233 |             
234 |             u2s_index = batch['user', 'u2s', 'server'].edge_index
235 |             user_index = u2s_index[0]
236 |             server_index = u2s_index[1]
237 |             u2s_path_loss = batch['user', 'u2s', 'server'].path_loss.squeeze().reshape((batch_size, -1))
238 |             u2s_path_loss_feat = batch['user', 'u2s', 'server'].edge_attr.squeeze().reshape((batch_size, -1))
239 |             user_tasksize = batch['user'].x[:, 0].reshape((batch_size, -1))
240 |             server_comp_resource = batch['server'].x[:, 0].reshape((batch_size, -1))
241 | 
242 |             task_sche, power_sche, comp_sche, sches, critic_value = model(u2s_path_loss_feat, user_tasksize, server_comp_resource, u2s_index)
243 |             total_time = compute_loss_nn(task_sche, power_sche, comp_sche, user_tasksize, server_comp_resource, u2s_path_loss, user_index, server_index)
244 |             
245 |             # u2s_index = batch['user', 'u2s', 'server'].edge_index
246 |             # u2s_path_loss = batch['user', 'u2s', 'server'].edge_attr.squeeze()
247 |             # user_tasksize = batch['user'].x[:, 0]
248 |             # server_comp_resource = batch['server'].x[:, 0]
249 |             # task_sche, power_sche, comp_sche, sches, critic_value = model(u2s_path_loss, user_tasksize, server_comp_resource, u2s_index)
250 |             # task_sche, power_sche, comp_sche, sches, critic_value = self.nnCritic_model(u2s_path_loss, u2s_index)
251 |             # total_time = compute_loss_nn(task_sche, power_sche, comp_sche, batch['server'].x[:, 0], batch['user', 'u2s', 'server'].path_loss, batch['user'].x[:, 0], u2s_index)
252 |             
253 |             sches_list.append(sches)
254 | 
255 |             critic_loss = critic_loss_func(total_time.detach(), critic_value)
256 |             critic_loss.backward()
257 | 
258 |             # if idx % batch_size == 0:
259 |             critic_optimizer.step()
260 |             critic_optimizer.zero_grad()
261 |             
262 |             sches_list = torch.vstack(sches_list)
263 |             value_loss = -model.critic_mlp(sches_list)
264 |             sches_list = []
265 |             value_loss.mean().backward()
266 |             actor_optimizer.step()
267 |             actor_optimizer.zero_grad()
268 |             sum_loss += total_time.item()
269 |             length += 1
270 |         
271 |         mean_loss = sum_loss/length
272 |         train_losses.append(mean_loss)
273 |         print('Time step: {} \t\t PCNetCritic training loss: {}'.format(step, mean_loss))
274 |     
275 |     return train_losses
276 | 
277 | 
278 | 
279 | def MMSE_test(loader):
280 |     mean_loss_list = []
281 |     task_sche_list = []
282 |     power_sche_list = []
283 |     comp_sche_list = []
284 |     for idx, batch in tqdm(enumerate(loader)):
285 |         compute_resource = batch['server'].x[:, 0]
286 |         path_loss = batch['user', 'u2s', 'server'].path_loss
287 |         edge_index = batch['user', 'u2s', 'server'].edge_index
288 |         task_size = batch['user'].x[:, 0]
289 |         
290 |         model = MMSE(input_shape=path_loss.shape, args=args).to(args.device)
291 |         lr = mse_lr
292 |         task_optimizer = torch.optim.SGD([{'params': model.task_allocation}], lr, momentum=0.9)
293 |         power_optimizer = torch.optim.SGD([{'params': model.power_allocation}], lr, momentum=0.9)
294 |         comp_optimizer = torch.optim.SGD([{'params': model.comp_allocation}], lr, momentum=0.9)
295 |         # optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
296 |         
297 |         for time_step in range(mse_epochs):
298 |             batch_loss = model(compute_resource, path_loss, task_size, edge_index)
299 |             batch_loss.backward()
300 |             # optimizer.step()
301 |             if time_step % 3 == 0:
302 |                 task_optimizer.step()
303 |                 task_optimizer.zero_grad()
304 |             elif time_step % 3 == 1:
305 |                 power_optimizer.step()
306 |                 power_optimizer.zero_grad()
307 |             elif time_step % 3 == 2:
308 |                 comp_optimizer.step()
309 |                 comp_optimizer.zero_grad()
310 |             # if time_step % 2 == 0:
311 |             #     task_optimizer.step()
312 |             # elif time_step % 2 == 1:
313 |             #     comp_optimizer.step()
314 |             mean_loss_list.append(batch_loss.item())
315 |             print('batch: {}, time_step: {}, loss : {}'.format(idx, time_step, batch_loss.item()))
316 |         print('batch: {} \t\t loss: {}'.format(idx, batch_loss.item()))
317 |         task_sche_list.append(model.task_allocation.data)
318 |         power_sche_list.append(model.power_allocation.data)
319 |         comp_sche_list.append(model.comp_allocation.data)
320 | 
321 |     return task_sche_list, power_sche_list, comp_sche_list
322 | 
323 | 
324 | def gnn_mse_supervised_train(layouts, model, task_target, power_target, comp_target):
325 | 
326 |     # task_target, power_target, comp_target = MMSE_test(layouts)
327 |     loss_func = torch.nn.MSELoss()
328 |     optimizer = torch.optim.Adam(model.parameters(), lr=hgnn_sv_lr)
329 |     policy_losses = []
330 |     for epoch in tqdm(range(Epochs)):
331 |         
332 |         # training
333 |         model.train()
334 |         loss_sum = 0
335 |         length = 0
336 |         for graph, tt, pt, ct in zip(layouts, task_target, power_target, comp_target):    # graph为一个batch
337 |             task_allocation, power_allocation, comp_allocation = model(graph.x_dict, graph.edge_index_dict, graph.edge_attr_dict)
338 |             
339 |             loss_batch = compute_loss(task_allocation, power_allocation, comp_allocation, graph['server'].x[:, 2], graph['user', 'u2s', 'server'].path_loss, graph['user'].x[:, 2], graph['user', 'u2s', 'server'].edge_index)
340 |             task_loss = loss_func(task_allocation, tt)
341 |             power_loss = loss_func(power_allocation, pt)
342 |             comp_loss = loss_func(comp_allocation, ct)
343 | 
344 |             loss = task_loss + power_loss + comp_loss
345 |             
346 |             optimizer.zero_grad()
347 |             loss.backward()
348 |             optimizer.step()
349 |             loss_sum += loss_batch.item()
350 |             length += 1
351 |         policy_loss = loss_sum/length
352 |         policy_losses.append(policy_loss)
353 |         print('epoch=={}, HgnnSupervised policy_loss=={}'.format(epoch, policy_loss))
354 |     
355 |     return policy_losses
356 | 
357 | def pcnet_mse_supervised_train(layouts, model, task_target, power_target, comp_target):
358 |     
359 |     loss_func = torch.nn.MSELoss()
360 |     optimizer = torch.optim.Adam(model.parameters(), lr=pcnet_sv_lr)
361 |     policy_losses = []
362 |     for epoch in tqdm(range(Epochs)):
363 |         
364 |         # training
365 |         model.train()
366 |         loss_sum = 0
367 |         length = 0
368 |         for graph, tt, pt, ct in zip(layouts, task_target, power_target, comp_target):    # graph为一个batch
369 |             
370 |             u2s_index = graph['user', 'u2s', 'server'].edge_index
371 |             u2s_path_loss = graph['user', 'u2s', 'server'].edge_attr.squeeze()
372 |             user_tasksize = graph['user'].x[:, 2]
373 |             server_comp_resource = graph['server'].x[:, 2]
374 |             task_sche, power_sche, comp_sche = model(u2s_path_loss, user_tasksize, server_comp_resource, u2s_index)
375 |             total_time = compute_loss(task_sche, power_sche, comp_sche, server_comp_resource, graph['user', 'u2s', 'server'].path_loss, user_tasksize, u2s_index)
376 |             
377 |             # task_allocation, power_allocation, comp_allocation = model(graph.x_dict, graph.edge_index_dict, graph.edge_attr_dict)
378 |             
379 |             # loss_batch = compute_loss(task_allocation, power_allocation, comp_allocation, graph['server'].x[:, 2], graph['user', 'u2s', 'server'].path_loss, graph['user'].x[:, 2], graph['user', 'u2s', 'server'].edge_index)
380 |             task_loss = loss_func(task_sche, tt)
381 |             power_loss = loss_func(power_sche, pt)
382 |             comp_loss = loss_func(comp_sche, ct)
383 | 
384 |             loss = task_loss + power_loss + comp_loss
385 | 
386 |             optimizer.zero_grad()
387 |             loss.backward()
388 |             optimizer.step()
389 |             loss_sum += total_time.item()
390 |             length += 1
391 |         policy_loss = loss_sum/length
392 | 
393 |         policy_losses.append(policy_loss)
394 |         print('epoch=={}, PcNetSupervised policy_loss=={}'.format(epoch, policy_loss))
395 |     
396 |     return policy_losses
397 | 
398 | def HGNN_train(model, train_loader):
399 |     policy_losses = []
400 |     optimizer = torch.optim.Adam(model.parameters(), lr=hgnn_lr)
401 |     optimizer_stepLR = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20,  gamma=0.9)
402 |     for time_step in tqdm(range(Epochs)):
403 |         
404 |         # training
405 |         model.train()
406 |         loss_sum = 0
407 |         length = 0
408 |         for graph in train_loader:    # graph为一个batch
409 |             task_allocation, power_allocation, comp_allocation = model(graph.x_dict, graph.edge_index_dict, graph.edge_attr_dict)
410 |             
411 |             loss_batch = compute_loss(task_allocation, power_allocation, comp_allocation, graph['server'].x[:, 0], graph['user', 'u2s', 'server'].path_loss, graph['user'].x[:, 0], graph['user', 'u2s', 'server'].edge_index)
412 |             optimizer.zero_grad()
413 |             loss_batch.backward()
414 |             optimizer.step()
415 |             loss_sum += loss_batch.item()
416 |             length += 1
417 |         optimizer_stepLR.step()
418 |         policy_loss = loss_sum/length
419 | 
420 |         policy_losses.append(policy_loss)
421 |         print('step=={}, policy_loss=={}'.format(time_step, policy_loss))
422 |     return policy_losses
423 | 
424 | 
425 | def NN_train(model, loader):
426 |     sum_loss = 0
427 |     length = 0
428 |     train_losses = []
429 |     test_losses = []
430 | 
431 |     optimizer = torch.optim.Adam(model.parameters(), lr=pcnet_lr)
432 |     schedular = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.98)
433 |     for step in tqdm(range(Epochs)):
434 |         model.train()
435 |         for idx, batch in enumerate(loader):
436 |             u2s_index = batch['user', 'u2s', 'server'].edge_index
437 |             user_index = u2s_index[0]
438 |             server_index = u2s_index[1]
439 |             u2s_path_loss = batch['user', 'u2s', 'server'].path_loss.squeeze().reshape((batch_size, -1))
440 |             u2s_path_loss_feat = batch['user', 'u2s', 'server'].edge_attr.squeeze().reshape((batch_size, -1))
441 |             user_tasksize = batch['user'].x[:, 0].reshape((batch_size, -1))
442 |             server_comp_resource = batch['server'].x[:, 0].reshape((batch_size, -1))
443 | 
444 |             task_sche, power_sche, comp_sche = model(u2s_path_loss_feat, user_tasksize, server_comp_resource, u2s_index)
445 |             loss = compute_loss_nn(task_sche, power_sche, comp_sche, user_tasksize, server_comp_resource, u2s_path_loss, user_index, server_index)
446 |             loss.backward()
447 |             # if idx % batch_size == 0:
448 |             optimizer.step()
449 |             optimizer.zero_grad()
450 |             sum_loss += loss.item()
451 |             length += 1
452 |         schedular.step()
453 |         
454 |         mean_loss = sum_loss/length
455 |         train_losses.append(mean_loss)
456 |         print('Time step: {} \t\t PCNet training loss: {}'.format(step, mean_loss))
457 |         
458 |     
459 |     return train_losses, test_losses
460 | 
461 | 
462 | 
463 | if __name__ == '__main__':
464 |     # 生成场景
465 |     np.random.seed(50)
466 |     torch.cuda.manual_seed_all(50)
467 |     num_layouts = 256
468 |     server_num = 15
469 |     batch_size = 16
470 |     Epochs = 2000
471 | 
472 |     hgnn_lr = 5e-3
473 |     pcnet_lr = 1e-5
474 |     
475 |     hgnn_sv_lr = 1e-4
476 |     pcnet_sv_lr = 5e-4
477 | 
478 |     actor_lr = 1e-4
479 |     critic_lr = 1e-4
480 | 
481 |     mse_epochs = 300
482 |     mse_lr = 0.05
483 | 
484 |     train_user_nums = np.random.randint(server_num*3, server_num*3+1, num_layouts)
485 |     train_server_nums = np.random.randint(server_num, server_num+1, num_layouts)
486 |     # test_user_nums = np.random.randint(server_num*3, server_num*3+1, args.test_layouts)
487 |     # test_server_nums = np.random.randint(server_num, server_num+1, args.test_layouts)
488 | 
489 |     # env_max_length = np.sqrt(server_num * 300)
490 | 
491 |     train_layouts = generate_layouts(train_user_nums, train_server_nums, args)
492 |     # test_layouts = generate_layouts(test_user_nums, test_server_nums, args)
493 | 
494 |     train_loader = DataLoader(train_layouts, batch_size=batch_size, shuffle=False)
495 |     # test_loader = DataLoader(test_layouts, batch_size=args.batch_size, shuffle=False)
496 | 
497 |     # single_task_target, single_power_target, single_comp_target = MMSE_test(train_layouts)
498 |     # batch_task_target, batch_power_target, batch_comp_target = MMSE_test(train_loader)
499 | 
500 |     # gnn_supervised_model = TaskLoad(args.num_layers, args.input_dim, args.hidden_dim, args.max_server_num, args.alpha).to(args.device)
501 |     gnn_model = TaskLoad(args.num_layers, args.input_dim, args.hidden_dim, args.max_server_num, args.alpha).to(args.device)
502 |     # gnn_model = torch.load('./TO_models/medium/medium_600.pt', map_location=args.device)
503 |     # hgnn_supervised_losses = gnn_mse_supervised_train(train_loader, gnn_supervised_model, batch_task_target, batch_power_target, batch_power_target)
504 |     # hgnn_losses = HGNN_train(gnn_model, train_layouts)
505 | 
506 |     
507 |     
508 |     
509 |     # PcNetCritic_model = PCNetCritic((server_num**2)*3+server_num*4, args.hidden_dim, (server_num**2)*3, args.alpha, args).to(args.device)
510 |     # HGnnCritic_model = GnnCritic(args.num_layers, args.input_dim, args.hidden_dim, (server_num**2)*3, args.alpha).to(args.device)
511 | 
512 |     # pcNet_critic_losses = NN_critic_train(PcNetCritic_model, train_loader, batch_size)
513 |     # HGNN_critic_losses= gnnCritic_train(HGnnCritic_model, train_layouts, batch_size)
514 | 
515 | 
516 |     
517 | 
518 |     pcnet = PCNet((server_num**2)*3+server_num*4, args.hidden_dim, (server_num**2)*3, args.alpha).to(args.device)
519 |     pcnet_supervised_model = PCNet((server_num**2)*3+server_num*4, args.hidden_dim, (server_num**2)*3, args.alpha).to(args.device)
520 | 
521 |     # pcnet_supervised_losses = pcnet_mse_supervised_train(train_layouts, pcnet_supervised_model, single_task_target, single_power_target, single_comp_target)
522 |     pcnet_losses = NN_train(pcnet, train_loader)
523 | 


--------------------------------------------------------------------------------
/layouts.py:
--------------------------------------------------------------------------------
  1 | import numpy as np
  2 | import torch
  3 | from torch_geometric.data import HeteroData
  4 | from arguments import args
  5 | import torch_geometric.transforms as T
  6 | 
  7 | 
  8 | 
  9 | def EuclideanDistances(a, b):
 10 |     '''
 11 |         Calculate the euclidean distances
 12 |     '''
 13 |     sq_a = a**2
 14 |     sum_sq_a = np.sum(sq_a, axis=1)
 15 |     sum_sq_a = np.expand_dims(sum_sq_a, axis=1)       # m->[m, 1]
 16 |     sq_b = b**2
 17 |     sum_sq_b = np.sum(sq_b, axis=1)  
 18 |     sum_sq_b = np.expand_dims(sum_sq_b, axis=0)        # n->[1, n]
 19 |     bt = b.T
 20 |     distance = np.sqrt(np.abs(sum_sq_a+sum_sq_b-2*np.matmul(a, bt)))
 21 |     return distance
 22 | 
 23 | 
 24 | def build_graph(user, server, s2u_idx, u2s_index, u2u_idx, s2s_idx, u2u_distance, u2s_path_loss, u2s_path_loss_feat, env_len):
 25 |     
 26 |     user_feat = user
 27 |     server_feat = server
 28 |     user_ones = np.zeros((len(user), 1))
 29 |     server_ones = np.zeros((len(server), 1))
 30 |     user_feat = np.concatenate((user_feat, user_ones), axis=1)
 31 |     server_feat = np.concatenate((server_feat, server_ones), axis=1)
 32 |     user_feat = torch.tensor(user_feat, dtype=torch.float).to(args.device)
 33 |     server_feat = torch.tensor(server_feat, dtype=torch.float).to(args.device)
 34 |     u2s_path_loss = torch.tensor(u2s_path_loss, dtype=torch.float).to(args.device)
 35 |     
 36 |     u2s_path_loss_feat = torch.tensor(u2s_path_loss_feat, dtype=torch.float).to(args.device)
 37 | 
 38 |     s2u_attr = u2s_path_loss_feat.reshape((-1, 1))
 39 |     u2s_attr = u2s_path_loss_feat.reshape((-1, 1))
 40 |     
 41 |     u2u_attr = torch.tensor(u2u_distance.reshape((-1, 1)), dtype=torch.float).to(args.device)
 42 | 
 43 | 
 44 |     s2u_idx = torch.tensor(s2u_idx, dtype=torch.long).to(args.device)
 45 |     u2s_index = torch.tensor(u2s_index, dtype=torch.long).to(args.device)
 46 |     u2u_idx = torch.tensor(u2u_idx, dtype=torch.long).to(args.device)
 47 |     s2s_idx = torch.tensor(s2s_idx, dtype=torch.long).to(args.device)
 48 | 
 49 |     data = HeteroData().to(args.device)
 50 | 
 51 |     data['env_len'].x = torch.tensor([env_len])
 52 | 
 53 | 
 54 |     data['user'].x = user_feat      # locations of users，size of task to offloading
 55 |     data['server'].x = server_feat
 56 |     data['user', 'u2u', 'user'].edge_index = u2u_idx
 57 |     data['server', 's2u', 'user'].edge_index = s2u_idx
 58 |     data['user', 'u2s', 'server'].edge_index = u2s_index
 59 | 
 60 |     data['user', 'u2s', 'server'].path_loss = u2s_path_loss_feat.reshape((-1, 1))
 61 | 
 62 |     data['server', 's2s', 'server'].edge_index = s2s_idx
 63 |     data['server', 's2u', 'user'].edge_attr = s2u_attr
 64 |     data['user', 'u2s', 'server'].edge_attr = u2s_attr
 65 |     data['user', 'u2u', 'user'].edge_attr = u2u_attr
 66 | 
 67 |     return data
 68 | 
 69 | def compute_path_losses(args, distances):
 70 |     carrier_f = (args.carrier_f_start+args.carrier_f_end)/2
 71 |     carrier_lam = 2.998e8 / carrier_f
 72 |     signal_cof = args.signal_cof
 73 |     path_losses = (signal_cof * carrier_lam) / (distances**2)
 74 | 
 75 |     return path_losses
 76 | 
 77 | def generate_layouts(user_nums, server_nums,args):
 78 |     
 79 |     graphs = []
 80 |     
 81 |     for idx in range(len(server_nums)):
 82 | 
 83 |         env_len = np.sqrt(server_nums[idx]*50)
 84 | 
 85 |         # 归一化位置, tasksize, computing resource
 86 |         user_idx_feat = np.random.random([user_nums[idx], 2])
 87 |         user_idx = user_idx_feat * env_len
 88 |         
 89 |         # user_idx_tasksize = np.random.random((user_nums[idx], 1))*(args.init_max_size-args.init_min_size) + args.init_min_size
 90 |         user_idx_tasksize = np.random.random((user_nums[idx], 1))
 91 |         user_idx_feat = user_idx_tasksize.copy()
 92 | 
 93 |         # server_idx_comp = np.random.random((server_nums[idx], 1))*(args.init_max_comp - args.init_min_comp) + args.init_min_comp
 94 |         server_idx_comp = np.random.random((server_nums[idx], 1))
 95 |         server_idx_feat = server_idx_comp.copy()
 96 | 
 97 |         server_idx_feat = np.random.random([server_nums[idx], 2])
 98 |         server_idx = server_idx_feat * env_len
 99 | 
100 | 
101 |         mask_ones = np.ones(server_nums[idx])
102 |         mask_idx = np.repeat(mask_ones[np.newaxis, :], repeats=user_nums[idx], axis=0)
103 | 
104 |         
105 |         user_num_idx = user_nums[idx]
106 |         # server_num_idx = server_nums[idx]
107 | 
108 |         # edge_index of users to users
109 |         index_src = np.arange(user_num_idx).repeat(repeats=user_num_idx)
110 |         index_dst = np.tile(np.arange(user_num_idx), reps=user_num_idx)
111 |         u2u_index = np.concatenate([index_src[np.newaxis, :], index_dst[np.newaxis, :]], axis=0)
112 | 
113 |         # edge_index of servers to users
114 |         index_s2u_dst, index_s2u_src = np.nonzero(mask_idx)
115 |         s2u_index = np.concatenate([index_s2u_src[np.newaxis, :], index_s2u_dst[np.newaxis, :]], axis=0)
116 | 
117 |         # edge_index of users to servers
118 |         index_u2s_src, index_u2s_dst = np.nonzero(mask_idx)
119 |         u2s_index = np.concatenate([index_u2s_src[np.newaxis, :], index_u2s_dst[np.newaxis, :]], axis=0)
120 | 
121 |         # edge_index of servers to servers
122 |         edge_mat = mask_ones[np.newaxis, :].repeat(repeats=server_nums[idx], axis=0)
123 |         index_s2s_src, index_s2s_dst = np.nonzero(edge_mat)
124 |         s2s_index = np.concatenate([index_s2s_src[np.newaxis, :], index_s2s_dst[np.newaxis, :]], axis=0)
125 | 
126 |         u2u_distances_idx = EuclideanDistances(user_idx, user_idx)
127 |         u2s_distances_idx = EuclideanDistances(user_idx, server_idx)
128 |         u2s_path_loss = compute_path_losses(args, u2s_distances_idx)
129 | 
130 | 
131 |         # generate normalized channels randomly
132 |         u2s_path_loss_feat = torch.rand((user_nums[idx], server_nums[idx]), device=args.device)
133 | 
134 |         # normalize the distances
135 |         distance_mean = np.mean(u2s_distances_idx)
136 |         distance_var = np.sqrt(np.mean(np.square(u2s_distances_idx-distance_mean)))
137 |         u2s_distances_idx = (u2s_distances_idx-distance_mean)/distance_var
138 |  
139 |         user_idx_feat = user_idx_tasksize
140 |         server_idx_feat = server_idx_comp
141 |         
142 |         graph = build_graph(user_idx_feat, server_idx_feat, s2u_index, u2s_index, u2u_index, s2s_index, u2u_distances_idx, u2s_path_loss, u2s_path_loss_feat, env_len)
143 |         graphs.append(graph)
144 |     
145 |     return graphs
146 | 


--------------------------------------------------------------------------------
/multi_scales_test.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | from arguments import args
  3 | import numpy as np
  4 | from torch_geometric.loader import DataLoader
  5 | from layouts import generate_layouts
  6 | from off_loading_models import TaskLoad, PCNet
  7 | from tqdm import tqdm
  8 | import random
  9 | from time import time
 10 | import pandas as pd
 11 | from GA_exp import Evolution
 12 | 
 13 | 
 14 | def compute_loss_nn(task_allocation, power_allocation, comp_allocation, task_size, compute_resource, path_losses, user_index, server_index):
 15 |     
 16 |     # task_size : vector N
 17 |     # task_allocation: mat pop_num x 3*M*N
 18 |     # index: vector 3*M*N
 19 |     
 20 |     epsilon = 1e-9
 21 |     extre = 1e-20
 22 |     server_index_first = server_index.reshape((batch_size, -1))[0]
 23 |     user_index_first = user_index.reshape((batch_size, -1))[0]
 24 |     # user_index = edge_index[0]      # s2u中源节点的索引
 25 |     # server_index = edge_index[1]    # s2u中目标节点的索引
 26 |     
 27 |     # power_allocation = torch.clamp(power_allocation, 1e-5, 1)
 28 |     pw_ini = power_allocation * path_losses    # mat pop_num x M*N
 29 |     # pw = torch.clamp(pw, 1e-5, 1)
 30 |     
 31 |     # 将信道状态过小的设置为0
 32 |     mask_pw = torch.where(pw_ini<args.pw_threshold)
 33 |     
 34 |     pw = pw_ini.clone()
 35 | 
 36 |     pw[mask_pw] = 0
 37 | 
 38 |     comp_allocation_clone = comp_allocation.clone()
 39 |     comp_allocation_clone[mask_pw] = 0
 40 |     comp_allocation_normed = torch.zeros((comp_allocation_clone.shape[0], comp_allocation_clone.shape[1], server_index_first[-1]+1), device=args.device)
 41 |     comp_allocation_normed.scatter_(2, server_index_first.repeat((batch_size, 1)).unsqueeze(2), comp_allocation_clone.unsqueeze(2))
 42 |     comp_allocation_normed = comp_allocation_normed.sum(1)[:, server_index_first]
 43 |     comp_allocation_normed = torch.div(comp_allocation_clone, comp_allocation_normed+extre)
 44 | 
 45 |     task_allocation_clone = task_allocation.clone()
 46 |     task_allocation_clone[mask_pw] = 0
 47 |     task_allocation_normed = torch.zeros((task_allocation_clone.shape[0], task_allocation_clone.shape[1], user_index_first[-1]+1), device=args.device)
 48 |     task_allocation_normed.scatter_(2, user_index_first.repeat((batch_size, 1)).unsqueeze(2), task_allocation_clone.unsqueeze(2))
 49 |     task_allocation_normed = task_allocation_normed.sum(1)[:, user_index_first]
 50 |     task_allocation_normed = torch.div(task_allocation_clone, task_allocation_normed+extre)
 51 | 
 52 |     # 计算速率
 53 |     pw_list = torch.zeros((pw.shape[0], pw.shape[1], server_index_first[-1]+1), device=args.device)   # mat pop_num x MN x N
 54 |     pw_list.scatter_(2, server_index_first.repeat((batch_size, 1)).unsqueeze(2), pw.unsqueeze(2))
 55 |     pws_list = pw_list.sum(1)[:, server_index_first]  # mat pop_num x MN
 56 | 
 57 |     interference = pws_list-pw
 58 |     rate = torch.log2(1+torch.div(pw, interference+epsilon))
 59 |     # rate = args.band_width * torch.log2(1+torch.div(pw, interference+epsilon))
 60 | 
 61 | 
 62 |     task_size = task_size[:, user_index_first]       # M*N    
 63 |     # task_size = task_size[user_index]*args.tasksize_cof       # 重复采样映射到边中
 64 |     tasks = task_size * task_allocation_normed   # mat pop_num x M*N
 65 | 
 66 |     compute_resource = compute_resource[:, server_index_first]
 67 |     # compute_resource = compute_resource[server_index]*args.comp_cof       # 
 68 | 
 69 |     comp = compute_resource * comp_allocation_normed
 70 | 
 71 | 
 72 |     # offloading_time = torch.div(tasks, rate+extre) * (args.tasksize_cof/args.band_width)
 73 |     offloading_time = torch.div(tasks, rate+extre)
 74 | 
 75 |     # compute_time = torch.div(tasks, comp+extre) * (args.tasksize_cof*args.cons_factor/args.comp_cof)
 76 |     compute_time = torch.div(tasks, comp+extre)
 77 | 
 78 |     time_loss = offloading_time + compute_time      # pop_num x MN
 79 |     assert torch.isnan(time_loss).sum()==0
 80 | 
 81 | 
 82 |     time_loss_list = torch.zeros((time_loss.shape[0], time_loss.shape[1], user_index_first[-1]+1), device=args.device)
 83 |     time_loss_list.scatter_(2, user_index_first.repeat((batch_size, 1)).unsqueeze(2), time_loss.unsqueeze(2))
 84 |     time_loss_list = time_loss_list.sum(1)      # pop_num x MN
 85 | 
 86 |     return time_loss_list.mean()
 87 | 
 88 | 
 89 | def compute_loss(task_allocation, power_allocation, comp_allocation, compute_resource, path_losses, task_size, user_index, server_index):
 90 |     
 91 |     epsilon = 1e-9
 92 |     extre = 1e-20
 93 |     # user_index = edge_index[0]      # s2u中源节点的索引
 94 |     # server_index = edge_index[1]    # s2u中目标节点的索引
 95 |     
 96 |     pw_ini = power_allocation.squeeze() * path_losses.squeeze()
 97 |     # pw小于阈值的对应power设为0
 98 |     pw = pw_ini.clone()
 99 |     # mask_pw = torch.where(pw_ini<args.pw_threshold)
100 |     # pw[mask_pw] = 0
101 | 
102 |     pw_user_list = torch.zeros((len(pw), user_index[-1]+1), device=args.device)
103 |     # pw_user_ini_list = torch.zeros((len(pw), user_index[-1]+1), device=args.device)
104 |     # pw_user_ini_list.scatter_(1, user_index.unsqueeze(1), pw_ini.unsqueeze(1))
105 |     pw_user_list.scatter_(1, user_index.unsqueeze(1), pw_ini.unsqueeze(1))
106 |     # 如果某一个user的发射功率均位于阈值以下
107 |     pw_masked = pw_user_list.clone()
108 |     pw_masked[torch.where(pw_masked < args.pw_threshold)] = 0
109 |     invalid_index = torch.where(pw_masked.sum(0)==0)[0]   # 是否有对所有server都低于阈值的
110 |     # assert len(invalid_index)==0
111 |     max_pw_index = pw_user_list[:, invalid_index].argmax(0) # 对所有server的pw都低于阈值的user 取信号最强的server
112 |     pw_masked[max_pw_index, invalid_index] = pw_user_list[max_pw_index, invalid_index]
113 |     pw = pw_masked.sum(1)
114 |     mask_pw = torch.where(pw==0)
115 | 
116 |     pw_list = torch.zeros((len(pw), server_index[-1]+1), device=args.device)
117 |     pw_list.scatter_(1, server_index.unsqueeze(1), pw.unsqueeze(1))
118 |     
119 |     pws_list = pw_list.sum(0)[server_index]
120 | 
121 | 
122 |     interference = pws_list-pw
123 |     rate = torch.log2(1+torch.div(pw, interference+epsilon))
124 | 
125 |     task_allocation_clone = task_allocation.clone().squeeze() + 1e-8
126 |     task_allocation_clone[mask_pw] = 0
127 |     task_allocation_normed = torch.zeros((len(task_allocation_clone), user_index[-1]+1), device=args.device)
128 |     task_allocation_normed.scatter_(1, user_index.unsqueeze(1), task_allocation_clone.unsqueeze(1))
129 |     # assert len(torch.where(task_allocation_normed.sum(0)==0)[0]) == 0
130 |     task_allocation_normed_2 = task_allocation_normed.sum(0)[user_index]
131 |     task_allocation_final = torch.div(task_allocation_clone, task_allocation_normed_2+extre)
132 |     # task_allocation_clone =  softmax(task_allocation_clone, user_index)
133 | 
134 |     task_size = task_size[user_index]
135 |     # task_size = task_size[user_index]*args.tasksize_cof       # 重复采样映射到边中
136 |     
137 |     tasks = task_size * task_allocation_final
138 | 
139 |     comp_allocation_clone = comp_allocation.clone().squeeze()
140 |     comp_allocation_clone[mask_pw] = 0
141 |     comp_allocation_normed = torch.zeros((len(comp_allocation_clone), server_index[-1]+1), device=args.device)
142 |     comp_allocation_normed.scatter_(1, server_index.unsqueeze(1), comp_allocation_clone.unsqueeze(1))
143 |     comp_allocation_normed_2 = comp_allocation_normed.sum(0)[server_index]
144 |     comp_allocation_final = torch.div(comp_allocation_clone, comp_allocation_normed_2+extre)
145 |     # compute_resource = compute_resource[server_index]*args.comp_cof       # 
146 | 
147 |     compute_resource = compute_resource[server_index]
148 |     comp = compute_resource * comp_allocation_final
149 | 
150 | 
151 |     # offloading_time = torch.div(tasks, rate+extre) * (args.tasksize_cof/args.band_width)
152 |     offloading_time = torch.div(tasks, rate+extre)
153 | 
154 |     # compute_time = torch.div(tasks, comp+extre) * (args.tasksize_cof*args.cons_factor/args.comp_cof)
155 |     compute_time = torch.div(tasks, comp+extre)
156 | 
157 |     # compute_time = torch.clamp(compute_time, -1, 3000)
158 |     # offloading_time = torch.clamp(offloading_time, -1, 3000)
159 | 
160 |     time_loss = offloading_time + compute_time
161 |     assert torch.isnan(time_loss).sum()==0
162 | 
163 |     time_loss_list = torch.zeros((len(time_loss), user_index[-1]+1), device=args.device)
164 |     time_loss_list.scatter_(1, user_index.unsqueeze(1), time_loss.unsqueeze(1))
165 |     time_loss_list = time_loss_list.sum(0)
166 | 
167 |     return time_loss_list.mean()
168 | 
169 | def setup_seed(seed):
170 |      torch.manual_seed(seed)
171 |      torch.cuda.manual_seed_all(seed)
172 |      np.random.seed(seed)
173 |      random.seed(seed)
174 |      torch.backends.cudnn.deterministic = True
175 |      np.random.seed(seed)
176 | 
177 | 
178 | 
179 | if __name__=='__main__':
180 |     
181 |     setup_seed(20)
182 |     hgnn_lr = 1e-5
183 |     pcnet_lr = 1e-4
184 |     batch_size = 1
185 |     train_steps = 300
186 |     # 生成测试数据
187 |     multi_scales = 'small'
188 |     if multi_scales == 'small':
189 |         min_server = 2
190 |         max_server = 7
191 |     elif multi_scales == 'medium':
192 |         min_server = 15
193 |         max_server = 20
194 |     elif multi_scales == 'large':
195 |         min_server = 25
196 |         max_server = 30
197 |     sample_nums = 30
198 |     layouts_list = []
199 |     server_num_list = list(range(min_server, max_server+1))
200 |     for num in server_num_list:
201 |         server_nums = np.ones(sample_nums, dtype=np.int8) * num
202 |         user_nums = server_nums * 3
203 |         layouts = generate_layouts(user_nums, server_nums, args)
204 |         layouts_list.append(layouts)
205 | 
206 | 
207 |     # test with multi_scales
208 | 
209 |     
210 | 
211 |     # inference directly
212 |         # loading pretrained_model
213 | 
214 | 
215 | 
216 |     # fine-tunning
217 | 
218 |     # GA iteration
219 | 
220 |     # PcNet inference directly
221 | 
222 |     # PcNet training
223 | 
224 |     # task offloading trainning process
225 |     for layouts, server_num in zip(layouts_list, server_num_list):
226 |         print('Now test server_num == {}'.format(server_num))
227 |         loaded_hgnn = torch.load('./TO_models/Pretrain/'+multi_scales+'/HGNN_2048_500.pt', map_location=args.device)
228 |         loaded_nn = torch.load('./TO_models/Pretrain/'+multi_scales+'/PcNet_2048_500.pt', map_location=args.device)
229 |         # inference directly
230 |         IDgnn_loss_list = []
231 |         IDgnn_time_list = []
232 |         IDNN_loss_list = []
233 |         IDNN_time_list = []
234 | 
235 | 
236 |         # print('######################################################Inference Directly Testing######################################################')
237 |         # for graph in tqdm(layouts):
238 |         #     # inference HGNN directly
239 |         #     start = time()
240 |         #     IDGnn_task_allocation, IDGnn_power_allocation, IDGnn_comp_allocation = loaded_hgnn(graph.x_dict, graph.edge_index_dict, graph.edge_attr_dict)
241 |         #     end = time()
242 |         #     user_index = graph['user', 'u2s', 'server'].edge_index[0]
243 |         #     server_index = graph['user', 'u2s', 'server'].edge_index[1]
244 |         #     loss_batch = compute_loss(IDGnn_task_allocation, IDGnn_power_allocation, IDGnn_comp_allocation, graph['server'].x[:, 0], graph['user', 'u2s', 'server'].path_loss, graph['user'].x[:, 0], user_index, server_index)
245 |         #     IDgnn_time_list.append(end-start)
246 |         #     IDgnn_loss_list.append(loss_batch.item())
247 | 
248 |         #     # inference PcNet directly
249 |         #     u2s_index = graph['user', 'u2s', 'server'].edge_index
250 |         #     user_index = u2s_index[0]
251 |         #     server_index = u2s_index[1]
252 |         #     u2s_path_loss = graph['user', 'u2s', 'server'].path_loss.squeeze().reshape((batch_size, -1))
253 |         #     u2s_path_loss_feat = graph['user', 'u2s', 'server'].edge_attr.squeeze().reshape((batch_size, -1))
254 |         #     user_tasksize = graph['user'].x[:, 0].reshape((batch_size, -1))
255 |         #     server_comp_resource = graph['server'].x[:, 0].reshape((batch_size, -1))
256 | 
257 |         #     user_index_cliped = user_index[torch.where(server_index < min_server)]
258 |         #     path_losses_feat_cliped = u2s_path_loss_feat.squeeze()[torch.where(server_index < min_server)]
259 |         #     server_index_cliped = server_index[torch.where(server_index < min_server)]
260 |         #     server_index_cliped = server_index_cliped[torch.where(user_index_cliped < 3*min_server)]
261 |         #     user_index_cliped = user_index_cliped[torch.where(user_index_cliped < 3*min_server)]
262 |         #     path_losses_feat_cliped = path_losses_feat_cliped[torch.where(user_index_cliped < 3*min_server)].reshape((batch_size, -1))
263 |         #     start = time()
264 |         #     # 前min_server规模的方案
265 |         #     IDNN_task_sche, IDNN_power_sche, IDNN_comp_sche = loaded_nn(path_losses_feat_cliped, user_tasksize[:, :3*min_server], server_comp_resource[:, :min_server], [user_index_cliped, server_index_cliped])
266 |         #     end = time()
267 |         #     task_sche= torch.rand(u2s_path_loss_feat.shape, device=args.device)
268 |         #     power_sche = torch.rand(u2s_path_loss_feat.shape, device=args.device)
269 |         #     comp_sche = torch.rand(u2s_path_loss_feat.shape, device=args.device)
270 |         #     task_sche[:, :3*min_server**2] = IDNN_task_sche
271 |         #     power_sche[:, :3*min_server**2] = IDNN_power_sche
272 |         #     comp_sche[:, :3*min_server**2] = IDNN_comp_sche
273 |         #     IDNN_time_list.append(end-start)
274 |         #     loss = compute_loss_nn(task_sche, power_sche, comp_sche, user_tasksize, server_comp_resource, u2s_path_loss, user_index, server_index)
275 |         #     IDNN_loss_list.append(loss.item())
276 | 
277 |         # ID_df = pd.DataFrame({
278 |         #                 'ID_HGNN_time':IDgnn_time_list,
279 |         #                 'ID_NN_time':IDNN_time_list, 
280 |         #                 'ID_HGNN_loss':IDgnn_loss_list,
281 |         #                 'ID_NN_loss':IDNN_loss_list})
282 |         
283 |         # ID_df.to_excel('./data/multi_scales_test/'+multi_scales+'/InferenceDirectly_{}.xlsx'.format(server_num))
284 | 
285 |         # print('######################################################Fine_Tuning--Training Testing######################################################')
286 |         # # for graph in layouts:
287 |         # FTgnn_loss_list = []
288 |         # FTnn_loss_list = []
289 |         # FTgnn_time_list = []
290 |         # FTnn_time_list = []
291 |         # for graph in tqdm(layouts):    # graph为一个batch
292 |         #     gnn_loss_idx = []
293 |         #     hgnn_model = loaded_hgnn
294 |         #     nn_model = PCNet((server_num**2)*3+server_num*4, args.hidden_dim, (server_num**2)*3, args.alpha).to(args.device)
295 |         #     hgnn_optimizer = torch.optim.Adam(hgnn_model.parameters(), lr=hgnn_lr)
296 |         #     nn_optimizer = torch.optim.Adam(nn_model.parameters(), lr=pcnet_lr)
297 |         #     # hgnn fine-tuning
298 |         #     gnn_time_idx= []
299 |             
300 |         #     for time_step in range(train_steps):
301 |         #         start = time()
302 |         #         task_allocation, power_allocation, comp_allocation = hgnn_model(graph.x_dict, graph.edge_index_dict, graph.edge_attr_dict)
303 |         #         edge_index = graph['user', 'u2s', 'server'].edge_index
304 |         #         user_index = edge_index[0]
305 |         #         server_index = edge_index[1]
306 |         #         loss_batch = compute_loss(task_allocation, power_allocation, comp_allocation, graph['server'].x[:, 0], graph['user', 'u2s', 'server'].path_loss, graph['user'].x[:, 0], user_index, server_index)
307 |         #         hgnn_optimizer.zero_grad()
308 |         #         loss_batch.backward()
309 |         #         hgnn_optimizer.step()
310 |         #         end = time()
311 |         #         gnn_time_idx.append(end-start)
312 |         #         gnn_loss_idx.append(loss_batch.item())
313 |         #     gnn_loss_idx = np.array(gnn_loss_idx)
314 |         #     gnn_time_idx = np.array(gnn_time_idx)
315 |         #     FTgnn_loss_list.append(gnn_loss_idx)
316 |         #     FTgnn_time_list.append(gnn_time_idx)
317 |             
318 |         #     # pcNet fine-tunning
319 |         #     u2s_index = graph['user', 'u2s', 'server'].edge_index
320 |         #     user_index = u2s_index[0]
321 |         #     server_index = u2s_index[1]
322 |         #     u2s_path_loss = graph['user', 'u2s', 'server'].path_loss.squeeze().reshape((batch_size, -1))
323 |         #     u2s_path_loss_feat = graph['user', 'u2s', 'server'].edge_attr.squeeze().reshape((batch_size, -1))
324 |         #     user_tasksize = graph['user'].x[:, 0].reshape((batch_size, -1))
325 |         #     server_comp_resource = graph['server'].x[:, 0].reshape((batch_size, -1))
326 | 
327 |         #     # 神经网络直接优化
328 |         #     nn_loss_idx = []
329 |         #     nn_time_idx = []
330 |         #     for time_step in range(train_steps):
331 |         #         start = time()
332 |         #         task_sche, power_sche, comp_sche = nn_model(u2s_path_loss_feat, user_tasksize, server_comp_resource, u2s_index)
333 |         #         loss = compute_loss_nn(task_sche, power_sche, comp_sche, user_tasksize, server_comp_resource, u2s_path_loss, user_index, server_index)
334 |         #         nn_optimizer.zero_grad()
335 |         #         loss.backward()
336 |         #         nn_optimizer.step()
337 |         #         end = time()
338 |         #         nn_time_idx.append(end-start)
339 |         #         nn_loss_idx.append(loss.item())
340 |         #     nn_loss_idx = np.array(nn_loss_idx)
341 |         #     nn_time_idx = np.array(nn_time_idx)
342 |         #     FTnn_loss_list.append(nn_loss_idx)
343 |         #     FTnn_time_list.append(nn_time_idx)
344 | 
345 |             
346 |         # FTgnn_loss_list = np.vstack(FTgnn_loss_list)
347 |         # FTgnn_loss_list = FTgnn_loss_list.mean(axis=0)     # N个样本求平均, 一个server_num
348 | 
349 |         # FTgnn_time_list = np.vstack(FTgnn_time_list)
350 |         # FTgnn_time_list = FTgnn_time_list.mean(axis=0)
351 |         
352 |         # FTnn_loss_list = np.vstack(FTnn_loss_list)
353 |         # FTnn_loss_list = FTnn_loss_list.mean(axis=0)
354 | 
355 |         # FTnn_time_list = np.vstack(FTnn_time_list)
356 |         # FTnn_time_list = FTnn_time_list.mean(axis=0)
357 | 
358 |         # FT_df = pd.DataFrame({
359 |         #                 'FT_HGNN_loss': FTgnn_loss_list,
360 |         #                 'FT_NN_loss': FTnn_loss_list,
361 |         #                 'FT_HGNN_time': FTgnn_time_list,
362 |         #                 'FT_NN_time': FTnn_time_list
363 |         # })
364 |         # FT_df.to_excel('./data/multi_scales_test/'+multi_scales+'/FineTunning_{}.xlsx'.format(server_num))
365 | 
366 |         # GA 测试
367 | 
368 |         print('######################################################GA Testing######################################################')
369 |         pop_num=200
370 |         epochs=300
371 | 
372 |         ga_loss_list = []
373 |         ga_time_list = []
374 |         for data in tqdm(layouts):
375 |             compute_resource = data['server'].x[:, 0].squeeze()
376 |             path_losses = data['user', 'u2s', 'server'].path_loss.squeeze()
377 |             task_size = data['user'].x[:, 0].squeeze()
378 |             edge_index = data['user', 'u2s', 'server'].edge_index
379 |             ga = Evolution(server_num, server_num*3, pop_num, epochs, compute_resource=compute_resource, path_losses=path_losses, task_size=task_size, edge_index=edge_index)
380 |             loss, ga_time = ga.evolution()
381 |             ga_time_list.append(ga_time)
382 |             ga_loss_list.append(loss.data.cpu().numpy())
383 |         ga_loss_list = np.vstack(ga_loss_list)
384 |         ga_loss_list = ga_loss_list.mean(axis=0)
385 |         ga_time_list = np.vstack(ga_time_list)
386 |         ga_time_list = ga_time_list.mean(axis=0)
387 |         
388 |         GA_df = pd.DataFrame({
389 |             'GA_loss': ga_loss_list,
390 |             'GA_time': ga_time_list
391 |         })
392 |         GA_df.to_excel('./data/multi_scales_test/'+multi_scales+'/GA_{}.xlsx'.format(server_num))
393 | 


--------------------------------------------------------------------------------
/off_loading_models.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn as nn
  3 | from torch_geometric.nn import MessagePassing, HeteroConv
  4 | from torch_geometric.utils import softmax
  5 | import torch.nn.functional as F
  6 | from typing import Optional
  7 | from torch import Tensor
  8 | import numpy as np
  9 | import tqdm
 10 | from arguments import args
 11 | 
 12 | 
 13 | 
 14 | class TaskLoad(nn.Module):
 15 |     def __init__(self, num_layers, input_dim, hidden_dim, output_dim, alpha) -> None:
 16 |         super(TaskLoad, self).__init__()
 17 |         self.user_encoder = nn.Sequential(
 18 |             nn.Linear(input_dim, hidden_dim),
 19 |             nn.ReLU(),
 20 |             nn.Linear(hidden_dim, hidden_dim),
 21 |             nn.ReLU()
 22 |             )
 23 |         self.server_encoder = nn.Sequential(
 24 |             nn.Linear(input_dim, hidden_dim),
 25 |             nn.ReLU(),
 26 |             nn.Linear(hidden_dim, hidden_dim),
 27 |             nn.ReLU()
 28 |         )
 29 |         self.convs = nn.ModuleList()
 30 | 
 31 |         for layer in range(num_layers):
 32 |             conv = HeteroConv({
 33 |             ('server', 's2u', 'user'): S2UGNN(input_dim, hidden_dim, output_dim, alpha),
 34 |             ('user', 'u2s', 'server'): U2SGNN(input_dim, hidden_dim, output_dim, alpha)
 35 |             })
 36 |             self.convs.append(conv)
 37 | 
 38 |         self.task_allocation_mlp = nn.Sequential(
 39 |             nn.Linear(hidden_dim, hidden_dim),
 40 |             nn.ReLU(),
 41 |             nn.Linear(hidden_dim, 1),
 42 |             nn.LeakyReLU(alpha)
 43 |         )
 44 | 
 45 |         self.power_alloc_mlp = nn.Sequential(
 46 |             nn.Linear(hidden_dim, hidden_dim),
 47 |             nn.ReLU(),
 48 |             nn.Linear(hidden_dim, 1),
 49 |             nn.LeakyReLU(alpha)
 50 |         )
 51 | 
 52 |         self.comp_power_alloc_mlp = nn.Sequential(
 53 |             nn.Linear(hidden_dim, hidden_dim),
 54 |             nn.ReLU(),
 55 |             nn.Linear(hidden_dim, 1),
 56 |             nn.LeakyReLU(alpha)
 57 |         )
 58 | 
 59 | 
 60 | 
 61 |     def forward(self, x_dict, edge_index_dict, edge_attr_dict):
 62 |         x_dict['user'] = self.user_encoder(x_dict['user'])
 63 |         x_dict['server'] = self.server_encoder(x_dict['server'])
 64 |         x_dict_0 = x_dict.copy()
 65 |         for conv in self.convs:
 66 |             x_dict = conv(x_dict, edge_index_dict, edge_attr_dict)
 67 |             # x_dict['user']  = x_dict['user'] + x_dict_0['user']
 68 |             # x_dict['server'] = x_dict['server'] + x_dict_0['server']
 69 |             # x_dict_0 = x_dict.copy()
 70 |         s2u_attr = self.convs[-1].convs['server__s2u__user'].message_attr
 71 |         u2s_attr = self.convs[-1].convs['user__u2s__server'].message_attr
 72 | 
 73 |         user_index = edge_index_dict['user', 'u2s', 'server'][0]
 74 |         # server_index = edge_index_dict['user', 'u2s', 'server'][1]
 75 |         server_index = edge_index_dict['server', 's2u', 'user'][0]
 76 | 
 77 |         task_allocation = self.task_allocation_mlp(u2s_attr)
 78 |         task_allocation = softmax(task_allocation, index=user_index)
 79 |         
 80 |         power_allocation = self.power_alloc_mlp(u2s_attr)
 81 |         power_allocation = softmax(power_allocation, index=user_index)
 82 |         
 83 |         comp_allocation = self.comp_power_alloc_mlp(s2u_attr)
 84 |         comp_allocation = softmax(comp_allocation, index=server_index)
 85 | 
 86 |         return task_allocation, power_allocation, comp_allocation
 87 | 
 88 | class U2SGNN(MessagePassing):
 89 |     def __init__(self, input_dim, hidden_dim, output_dim, alpha, aggr: Optional[str] = "add", flow: str = "source_to_target", node_dim: int = -2, decomposed_layers: int = 1):
 90 |         super(U2SGNN, self).__init__()
 91 |         self.linear1 = nn.Linear(hidden_dim, hidden_dim)
 92 |         self.linear2 = nn.Linear(hidden_dim, hidden_dim)
 93 |         self.message_mlp = nn.Sequential(
 94 |             nn.Linear(hidden_dim+1, hidden_dim),
 95 |             nn.ReLU(),
 96 |         )
 97 |         self.update_lin = nn.Linear(2*hidden_dim, hidden_dim)
 98 |         self.att = nn.Sequential(
 99 |             nn.Linear(2*hidden_dim, hidden_dim),
100 |             nn.ReLU(),
101 |             nn.Linear(hidden_dim, 1),
102 |             nn.LeakyReLU(alpha)
103 |         )
104 | 
105 |         self.relation_mlp = nn.Sequential(
106 |             nn.Linear(2*hidden_dim+1, hidden_dim),
107 |             nn.ReLU(),
108 |             nn.Linear(hidden_dim, hidden_dim),
109 |             nn.ReLU()
110 |         )
111 | 
112 |         # self.task_alloc_mlp = nn.Sequential(
113 |         #     nn.Linear(hidden_dim, hidden_dim),
114 |         #     nn.ReLU(),
115 |         #     nn.Linear(hidden_dim, hidden_dim),
116 |         #     nn.LeakyReLU(alpha)
117 |         # )
118 | 
119 |         self.Wq = nn.Linear(hidden_dim, hidden_dim)
120 |         self.Wr = nn.Linear(hidden_dim, hidden_dim)
121 |     def forward(self, x, edge_index, edge_attr):
122 |         x_src = x[0]      # user
123 |         x_dst = x[1]       # server
124 |         # x_src = F.relu(self.linear1(x_src))
125 |         # x_dst = F.relu(self.linear2(x_dst))
126 |         return self.propagate(x=(x_src, x_dst), edge_index=edge_index, edge_attr=edge_attr)
127 |     def message(self, x_i, x_j, edge_index, edge_attr) -> Tensor:
128 |         # 消息传播机制
129 | 
130 |         # message_mlp计算用户到server的信息传递
131 |         tmp = torch.cat([x_i, edge_attr], dim=1)
132 |         # 计算注意力
133 |         att_weight = self.att(torch.cat([self.Wq(x_i), self.Wr(x_j)], dim=1))
134 |         att_weight = softmax(att_weight, index=edge_index[0], dim=0)        # 根据user进行softmax
135 | 
136 |         outputs = self.message_mlp(tmp)
137 |         outputs = att_weight*outputs        # 注意力
138 | 
139 |         # 将注意力特征与边特征结合得到user到server的关系特征
140 |         tmp = torch.cat([x_j, outputs, edge_attr], dim=-1)       # user的特征，server到user的关系特征与边特征
141 |         self.message_attr = self.relation_mlp(tmp)
142 | 
143 |         # # 关系特征mlp, softmax后得到user对server的分配结果
144 |         # self.task_allocation = self.task_alloc_mlp(message_attr)        # 每个user到server的分配信息
145 | 
146 |         return outputs
147 |     def update(self, aggr_out, x) -> Tensor:
148 | 
149 |         output = F.relu(self.update_lin(torch.column_stack([aggr_out, x[1]])))
150 |         
151 |         return output
152 | 
153 | 
154 | 
155 | class S2UGNN(MessagePassing):
156 |     def __init__(self, input_dim, hidden_dim, output_dim, alpha, aggr: Optional[str] = "add", flow: str = "source_to_target", node_dim: int = -2, decomposed_layers: int = 1):
157 |         super(S2UGNN, self).__init__()
158 |         self.linear1 = nn.Linear(hidden_dim, hidden_dim)
159 |         self.linear2 = nn.Linear(hidden_dim, hidden_dim)
160 |         self.message_mlp = nn.Sequential(
161 |             nn.Linear(hidden_dim+1, hidden_dim),
162 |             nn.ReLU(),
163 |             nn.Linear(hidden_dim, hidden_dim),
164 |             nn.ReLU()
165 |         )
166 |         self.update_lin = nn.Linear(2*hidden_dim, hidden_dim)
167 |         self.att = nn.Sequential(
168 |             nn.Linear(2*hidden_dim, 1),
169 |             nn.LeakyReLU(alpha)
170 |         )
171 | 
172 |         self.relation_mlp = nn.Sequential(
173 |             nn.Linear(2*hidden_dim+1, hidden_dim),
174 |             nn.ReLU(),
175 |             nn.Linear(hidden_dim, hidden_dim),
176 |             nn.ReLU()
177 |         )
178 |         # self.comp_power_alloc_mlp = nn.Sequential(
179 |         #     nn.Linear(hidden_dim, hidden_dim),
180 |         #     nn.ReLU(),
181 |         #     nn.Linear(hidden_dim, 1),
182 |         #     nn.LeakyReLU(alpha)
183 |         # )
184 | 
185 |         self.Wq = nn.Linear(hidden_dim, hidden_dim)
186 |         self.Wr = nn.Linear(hidden_dim, hidden_dim)
187 |     def forward(self, x, edge_index, edge_attr):
188 |         x_src = x[0]
189 |         x_dst = x[1]
190 |         # x_src = F.relu(self.linear1(x_src))
191 |         # x_dst = F.relu(self.linear2(x_dst))
192 |         return self.propagate(x=(x_src, x_dst), edge_index=edge_index, edge_attr=edge_attr)
193 |     def message(self, x_i, x_j, edge_index, edge_attr) -> Tensor:
194 |         # 消息传播机制
195 | 
196 |         # message_mlp计算server到user的信息传递
197 |         tmp = torch.cat([x_i, edge_attr], dim=1)
198 |         # 计算注意力
199 |         att_weight = self.att(torch.cat([self.Wq(x_i), self.Wr(x_j)], dim=1))
200 |         att_weight = softmax(att_weight, index=edge_index[0], dim=0)
201 | 
202 |         outputs = self.message_mlp(tmp)
203 |         outputs = att_weight*outputs
204 | 
205 |         # 将注意力特征与边特征结合得到user到server的关系特征
206 |         tmp = torch.cat([x_j, outputs, edge_attr], dim=-1)       # user的特征，server到user的关系特征与边特征
207 |         self.message_attr = self.relation_mlp(tmp)
208 | 
209 |         return outputs
210 |     def update(self, aggr_out, x) -> Tensor:
211 |         output = F.relu(self.update_lin(torch.column_stack([aggr_out, x[1]])))
212 |         return output
213 | 
214 | 
215 | '''
216 | class TaskLoad(nn.Module):
217 |     def __init__(self, input_dim, hidden_dim, output_dim, alpha) -> None:
218 |         super(TaskLoad, self).__init__()
219 |         self.user_linear = nn.Linear(input_dim, hidden_dim)
220 |         self.server_linear = nn.Linear(input_dim-1, hidden_dim)
221 |         
222 |         self.conv1 = HeteroConv({
223 |             ('server', 's2u' , 'user'): S2UGNN(input_dim, hidden_dim, output_dim, alpha)
224 |         })
225 |         self.conv2 = HeteroConv({
226 |             ('user', 'u2u', 'user'): U2UGNN(hidden_dim, hidden_dim, output_dim, alpha)
227 |         })
228 |         self.conv3 = HeteroConv({       # user到server，user根据server获得权值特征，输出user到server的边特征(分配方案)
229 |             ('user', 'u2s', 'server'): S2UGNN_U_Allocation(input_dim, hidden_dim, output_dim, alpha)
230 |         })
231 |         self.conv4 = HeteroConv({       # server到user，进行算力资源分配
232 |             ('user', 'u2s', 'server'):  S2UGNN_S_Allocation(input_dim, hidden_dim, output_dim, alpha)
233 |         })
234 | 
235 |     def forward(self, x_dict, edge_index_dict, edge_attr_dict):
236 |         # s2u，user获取所有server的信息
237 |         x_dict_1 = self.conv1(x_dict, edge_index_dict, edge_attr_dict)
238 |         x_dict['user'] = x_dict_1['user']
239 |         # u2u, user间共享位置信息
240 |         x_dict_2 = self.conv2(x_dict, edge_index_dict, edge_attr_dict)
241 |         x_dict['user'] = x_dict_2['user']
242 |         # u2s, user将特征与server进行注意力获取，得到卸载分配与功率分配方案
243 |         x_dict_3 = self.conv3(x_dict, edge_index_dict, edge_attr_dict)
244 |         task_allocation = self.conv3.convs['user__u2s__server'].task_allocation
245 |         power_allocation = self.conv3.convs['user__u2s__server'].power_allocation
246 |         # 将卸载分配与功率分配方案作为边特征输入到s2u
247 |         edge_attr_dict['user', 'u2s', 'server'] = torch.cat([edge_attr_dict['user', 'u2s', 'server'], task_allocation, power_allocation], dim=-1)
248 |         x_dict_4 = self.conv4(x_dict, edge_index_dict, edge_attr_dict)
249 |         comp_allocation = self.conv4.convs['user__u2s__server'].comp_power_allocation
250 | 
251 |         assert torch.isnan(task_allocation).sum()==0
252 |         assert torch.isnan(power_allocation).sum()==0
253 |         assert torch.isnan(comp_allocation).sum()==0
254 | 
255 |         return task_allocation, power_allocation, comp_allocation
256 | 
257 | 
258 | 
259 | 
260 | 
261 | 
262 | 
263 | 
264 | 
265 | class S2UGNN_U_Allocation(MessagePassing):
266 |     def __init__(self, input_dim, hidden_dim, output_dim, alpha, aggr: Optional[str] = "add", flow: str = "source_to_target", node_dim: int = -2, decomposed_layers: int = 1):
267 |         super().__init__()
268 |         self.linear1 = nn.Linear(hidden_dim, hidden_dim)
269 |         self.linear2 = nn.Linear(input_dim, hidden_dim)
270 | 
271 |         self.att = nn.Sequential(
272 |             nn.Linear(2*hidden_dim, 1),
273 |             nn.LeakyReLU(alpha)
274 |         )
275 | 
276 |         self.message_mlp = nn.Sequential(
277 |             nn.Linear(hidden_dim+1, hidden_dim),
278 |             nn.ReLU(),
279 |             nn.Linear(hidden_dim, hidden_dim),
280 |             nn.ReLU()
281 |         )
282 |         self.relation_mlp = nn.Sequential(
283 |             nn.Linear(2*hidden_dim+1, hidden_dim),
284 |             nn.ReLU()
285 |         )
286 |         self.task_alloc_mlp = nn.Sequential(
287 |             nn.Linear(hidden_dim, hidden_dim),
288 |             nn.ReLU(),
289 |             nn.Linear(hidden_dim, 1),
290 |             nn.LeakyReLU(alpha)
291 |         )
292 |         self.power_alloc_mlp = nn.Sequential(
293 |             nn.Linear(hidden_dim, hidden_dim),
294 |             nn.ReLU(),
295 |             nn.Linear(hidden_dim, 1),
296 |             nn.LeakyReLU(alpha)
297 |         )
298 |         self.Wq = nn.Linear(hidden_dim, hidden_dim)
299 |         self.Wr = nn.Linear(hidden_dim, hidden_dim)
300 |     def forward(self, x, edge_index, edge_attr):
301 |         user = x[0]
302 |         server = x[1]
303 |         user = F.relu(self.linear1(user))
304 |         server = F.relu(self.linear2(server))
305 |         return self.propagate(x=(user, server), edge_index=edge_index, edge_attr=edge_attr)
306 |     def message(self, x_i, x_j, edge_index, edge_attr) -> Tensor:
307 |         src_tmp = torch.cat([x_i, edge_attr], dim=1)
308 |         
309 |         # 计算server对user的注意力特征
310 |         att_weight = self.att(torch.cat([self.Wq(x_i), self.Wr(x_j)], dim=1))
311 |         att_weight = softmax(att_weight, index=edge_index[0])
312 | 
313 |         outputs = self.message_mlp(src_tmp)
314 | 
315 |         outputs = att_weight * outputs      # user到server的边特征，与edge_index维度相同
316 | 
317 |         # 将注意力特征与边特征结合得到user到server的关系特征
318 |         tmp = torch.cat([x_j, outputs, edge_attr], dim=-1)       # user的特征，server到user的关系特征与边特征
319 |         message_attr = self.relation_mlp(tmp)
320 | 
321 |         # 关系特征mlp, softmax后得到user对server的分配结果
322 |         self.task_allocation = self.task_alloc_mlp(message_attr)
323 |         self.power_allocation = self.power_alloc_mlp(message_attr)
324 | 
325 |         self.task_allocation = softmax(self.task_allocation, index=edge_index[0], dim=0)
326 |         self.power_allocation = softmax(self.power_allocation, index=edge_index[0], dim=0)
327 | 
328 | 
329 |         return super().message(x_j)
330 | 
331 | 
332 | class S2UGNN_S_Allocation(MessagePassing):
333 |     def __init__(self, input_dim, hidden_dim, output_dim, alpha, aggr: Optional[str] = "add", flow: str = "source_to_target", node_dim: int = -2, decomposed_layers: int = 1):
334 |         super().__init__()
335 |         self.linear1 = nn.Linear(hidden_dim, hidden_dim)
336 |         self.linear2 = nn.Linear(input_dim, hidden_dim)
337 |         self.message_mlp = nn.Sequential(
338 |             nn.Linear(hidden_dim+3, hidden_dim),
339 |             nn.ReLU()
340 |         )
341 |         self.att = nn.Sequential(
342 |             nn.Linear(2*hidden_dim, 1),
343 |             nn.LeakyReLU(alpha)
344 |         )
345 |         self.relation_mlp = nn.Sequential(
346 |             nn.Linear(2*hidden_dim+3, hidden_dim),
347 |             nn.ReLU(),
348 |             nn.Linear(hidden_dim, hidden_dim),
349 |             nn.ReLU()
350 |         )
351 |         self.comp_power_alloc_mlp = nn.Sequential(
352 |             nn.Linear(hidden_dim, hidden_dim),
353 |             nn.ReLU(),
354 |             nn.Linear(hidden_dim, 1),
355 |             nn.LeakyReLU(alpha)
356 |         )
357 |         self.Wq = nn.Linear(hidden_dim, hidden_dim)
358 |         self.Wr = nn.Linear(hidden_dim, hidden_dim)
359 |     def forward(self, x, edge_index, edge_attr):   # edge_attr为功率分配与task offloading分配结果，均为softmax比例
360 |         x_src = x[0]
361 |         x_dst = x[1]
362 |         self.tasksize = x_dst[:, 2]
363 |         x_src = F.relu(self.linear1(x_src))
364 |         x_dst = F.relu(self.linear2(x_dst))
365 |         return self.propagate(x=(x_src, x_dst), edge_index=edge_index, edge_attr=edge_attr)
366 |     def message(self, x_i, x_j, edge_index, edge_attr) -> Tensor:
367 |         src_tmp = torch.cat([x_j, edge_attr], dim=1)
368 |         
369 |         # 计算server对user的注意力特征
370 |         att_weight = self.att(torch.cat([self.Wq(x_i), self.Wr(x_j)], dim=1))
371 |         att_weight = softmax(att_weight, index=edge_index[1], dim=0)
372 | 
373 |         outputs = self.message_mlp(src_tmp)
374 | 
375 |         outputs = att_weight * outputs      # user到server的边特征，与edge_index维度相同
376 | 
377 |         # 将注意力特征与边特征结合得到user到server的关系特征
378 |         tmp = torch.cat([x_i, outputs, edge_attr], dim=-1)       # server的特征，server到user的关系特征与边特征
379 |         message_attr = self.relation_mlp(tmp)        # 
380 | 
381 |         # 关系特征mlp, softmax后得到server对user的算力分配结果
382 |         self.comp_power_allocation = self.comp_power_alloc_mlp(message_attr)
383 |         self.comp_power_allocation = softmax(self.comp_power_allocation, index=edge_index[1], dim=0)
384 | 
385 |         return super().message(x_j)
386 |         
387 |    
388 | 
389 | 
390 | class S2UGNN(MessagePassing):
391 |     def __init__(self, input_dim, hidden_dim, output_dim, alpha, aggr: Optional[str] = "add", flow: str = "source_to_target", node_dim: int = -2, decomposed_layers: int = 1):
392 |         super(S2UGNN, self).__init__()
393 |         self.linear1 = nn.Linear(input_dim, hidden_dim)
394 |         self.linear2 = nn.Linear(input_dim, hidden_dim)
395 |         self.message_mlp = nn.Sequential(
396 |             nn.Linear(hidden_dim+1, hidden_dim),
397 |             nn.ReLU(),
398 |         )
399 |         self.update_lin = nn.Linear(2*hidden_dim, hidden_dim)
400 |         self.att = nn.Sequential(
401 |             nn.Linear(2*hidden_dim, 1),
402 |             nn.LeakyReLU(alpha)
403 |         )
404 |         self.Wq = nn.Linear(hidden_dim, hidden_dim)
405 |         self.Wr = nn.Linear(hidden_dim, hidden_dim)
406 |     def forward(self, x, edge_index, edge_attr):
407 |         x_src = x[0]
408 |         x_dst = x[1]
409 |         x_src = F.relu(self.linear1(x_src))
410 |         x_dst = F.relu(self.linear2(x_dst))
411 |         return self.propagate(x=(x_src, x_dst), edge_index=edge_index, edge_attr=edge_attr)
412 |     def message(self, x_i, x_j, edge_index, edge_attr) -> Tensor:
413 |         # 消息传播机制
414 | 
415 |         # message_mlp计算用户到uav的信息传递
416 |         tmp = torch.cat([x_i, edge_attr], dim=1)
417 |         # 计算注意力
418 |         att_weight = self.att(torch.cat([self.Wq(x_i), self.Wr(x_j)], dim=1))
419 |         att_weight = softmax(att_weight, index=edge_index[0], dim=0)
420 | 
421 |         outputs = self.message_mlp(tmp)
422 |         outputs = att_weight*outputs
423 |         return outputs
424 |     def update(self, aggr_out, x) -> Tensor:
425 |         output = F.relu(self.update_lin(torch.column_stack([aggr_out, x[1]])))
426 |         return output
427 | 
428 | class U2UGNN(MessagePassing):
429 |     def __init__(self, input_dim, hidden_dim, ouput_dim, alpha, aggr: Optional[str] = "add", flow: str = "source_to_target", node_dim: int = -2, decomposed_layers: int = 1):
430 |         super().__init__()
431 |         self.linear1 = nn.Linear(input_dim, hidden_dim)
432 |         # self.allocation_linear = nn.Linear(hidden_dim, ouput_dim)
433 |         # self.power_linear = nn.Linear(hidden_dim, ouput_dim)
434 |         self.att = nn.Sequential(
435 |             nn.Linear(2*hidden_dim, 1),
436 |             nn.LeakyReLU(alpha)
437 |         )
438 |         self.message_mlp = nn.Sequential(
439 |             nn.Linear(2*hidden_dim, hidden_dim),
440 |             nn.ReLU(),
441 |         )
442 |         # self.update_mlp = nn.Sequential(
443 |         #     nn.Linear(2*hidden_dim, hidden_dim),
444 |         #     nn.ReLU(),
445 |         # )
446 |         self.update_lin = nn.Linear(2*hidden_dim, hidden_dim)
447 |         self.Wq = nn.Linear(hidden_dim, hidden_dim)
448 |         self.Wr = nn.Linear(hidden_dim, hidden_dim)
449 |     def forward(self, x, edge_index, edge_attr):
450 |         x = F.relu(self.linear1(x))
451 |         # self.user_num = edge_index[0][-1] + 1
452 |         return self.propagate(edge_index=edge_index, x=x)
453 |     def message(self, x_i, x_j, edge_index) -> torch.Tensor:
454 |         # 消息传播机制
455 | 
456 |         # message_mlp计算用户到uav的信息传递
457 |         self.outputs = self.message_mlp(torch.column_stack([x_i, x_j]))
458 | 
459 |         # 计算注意力
460 |         self.att_weight = self.att(torch.cat([self.Wq(x_i), self.Wr(x_j)], dim=1))
461 |         self.att_weight = softmax(self.att_weight, index=edge_index[0], dim=0)
462 |         self.outputs = self.att_weight * self.outputs
463 |         return self.outputs
464 | 
465 |     def aggregate(self, inputs: Tensor, index: Tensor, ptr: Optional[Tensor] = None, dim_size: Optional[int] = None) -> Tensor:
466 |         # 消息聚合
467 |         outputs = super().aggregate(inputs, index, ptr, dim_size)       # 进行aggr操作
468 |         return outputs
469 | 
470 |     def update(self, aggr_out, x):
471 | 
472 |         output = F.relu(self.update_lin(torch.column_stack([aggr_out, x])))
473 |         # 映射到01之间
474 |         # power = F.softmax(self.power_linear(x))
475 |         # allocation = F.softmax(self.allocation_linear(x))
476 | 
477 |         # ouput = torch.cat([allocation, power], dim=1)
478 | 
479 |         return output
480 | 
481 | '''
482 | 
483 | 
484 | class GnnCritic(nn.Module):
485 |     def __init__(self, num_layers, input_dim, hidden_dim, output_dim, alpha) -> None:
486 |         super(GnnCritic, self).__init__()
487 |         # output_dim: edge_index的数量 --> server_num x user_num
488 |         self.gnn = TaskLoad(num_layers, input_dim, hidden_dim, output_dim, alpha)
489 | 
490 |         self.critic_mlp = nn.Sequential(
491 |             nn.Linear(output_dim*3, hidden_dim),
492 |             nn.ReLU(),
493 |             nn.Linear(hidden_dim, hidden_dim),
494 |             nn.ReLU(),
495 |             nn.Linear(hidden_dim, 1),
496 |             nn.LeakyReLU(alpha)
497 |         )
498 | 
499 |     def forward(self, x_dict, edge_index_dict, edge_attr_dict):
500 |         # s2u，user获取所有server的信息
501 |         x_dict_1 = self.gnn.conv1(x_dict, edge_index_dict, edge_attr_dict)
502 |         x_dict['user'] = x_dict_1['user']
503 |         # u2u, user间共享位置信息
504 |         x_dict_2 = self.gnn.conv2(x_dict, edge_index_dict, edge_attr_dict)
505 |         x_dict['user'] = x_dict_2['user']
506 |         # u2s, user将特征与server进行注意力获取，得到卸载分配与功率分配方案
507 |         x_dict_3 = self.gnn.conv3(x_dict, edge_index_dict, edge_attr_dict)
508 |         task_allocation = self.gnn.conv3.convs['user__u2s__server'].task_allocation
509 |         power_allocation = self.gnn.conv3.convs['user__u2s__server'].power_allocation
510 |         # 将卸载分配与功率分配方案作为边特征输入到s2u
511 |         edge_attr_dict['user', 'u2s', 'server'] = torch.cat([edge_attr_dict['user', 'u2s', 'server'], task_allocation, power_allocation], dim=-1)
512 |         x_dict_4 = self.gnn.conv4(x_dict, edge_index_dict, edge_attr_dict)
513 |         comp_allocation = self.gnn.conv4.convs['user__u2s__server'].comp_power_allocation
514 | 
515 |         sches = torch.cat([task_allocation, power_allocation, comp_allocation], dim=0).squeeze()
516 |         critic_value = self.critic_mlp(sches.detach())
517 | 
518 |         return task_allocation, power_allocation, comp_allocation, sches, critic_value
519 | 
520 | 
521 | 
522 | 
523 | 
524 | 
525 | class PCNet(nn.Module):
526 |     def __init__(self, input_dim, hidden_dim, output_dim, alpha) -> None:
527 |         super(PCNet, self).__init__()
528 |         
529 |         self.output_dim = output_dim
530 | 
531 |         self.encoder = nn.Sequential(
532 |             nn.Linear(input_dim, hidden_dim),
533 |             nn.ReLU(),
534 |             nn.Linear(hidden_dim, hidden_dim),
535 |             nn.ReLU()
536 |         )
537 | 
538 |         self.sche_mlp = nn.Sequential(
539 |             nn.Linear(hidden_dim, hidden_dim),
540 |             nn.ReLU(),
541 |             nn.Linear(hidden_dim, 3*output_dim),
542 |             # nn.LeakyReLU(alpha)
543 |             # nn.Sigmoid()
544 |             # nn.Tanh()
545 |         )
546 | 
547 |         
548 |         
549 | 
550 |     def forward(self, u2s_path_loss, user_tasksize, server_comp_resource, edge_index):
551 |         '''
552 |             input: path_losses of users to servers
553 |             output: power_sche, task_sche, comp_sche
554 |         '''
555 |         feats = torch.cat([u2s_path_loss, user_tasksize, server_comp_resource], dim=1)
556 |         batch_num = feats.shape[0]
557 |         
558 |         embeddings = self.encoder(feats)
559 | 
560 |         sches = self.sche_mlp(embeddings)
561 |         task_sche = sches[:, :self.output_dim]
562 |         task_sche = softmax(task_sche.reshape((-1, 1)), index=edge_index[0])
563 |         task_sche = task_sche.reshape((batch_num, -1))
564 | 
565 |         power_sche = sches[:, self.output_dim:2*self.output_dim]
566 |         power_sche = softmax(power_sche.reshape((-1, 1)), edge_index[0])
567 |         power_sche = power_sche.reshape((batch_num, -1))
568 | 
569 |         comp_sche = sches[:, 2*self.output_dim:]
570 |         comp_sche = softmax(comp_sche.reshape((-1, 1)), index=edge_index[1])
571 |         comp_sche = comp_sche.reshape((batch_num, -1))
572 | 
573 |         # power_sche = self.power_mlp(embeddings)
574 |         # power_sche = softmax(power_sche, edge_index[0])
575 |         # task_sche = self.task_mlp(embeddings)
576 |         # task_sche = softmax(task_sche, index=edge_index[0])
577 |         # comp_sche = self.comp_mlp(embeddings)
578 |         # comp_sche = softmax(comp_sche, index=edge_index[1])
579 | 
580 |         return task_sche, power_sche, comp_sche
581 | 
582 | 
583 | class PCNetCritic(nn.Module):
584 |     def __init__(self, input_dim, hidden_dim, output_dim, alpha, args) -> None:
585 |         super(PCNetCritic, self).__init__()
586 |         
587 |         self.output_dim = output_dim
588 | 
589 |         self.encoder = nn.Sequential(
590 |             nn.Linear(input_dim, hidden_dim),
591 |             nn.ReLU(),
592 |             nn.Linear(hidden_dim, hidden_dim),
593 |             nn.ReLU()
594 |         )
595 | 
596 |         self.sche_mlp = nn.Sequential(
597 |             nn.Linear(hidden_dim, hidden_dim),
598 |             nn.ReLU(),
599 |             nn.Linear(hidden_dim, 3*output_dim),
600 |             nn.Tanh()
601 |         )
602 |         
603 | 
604 |         self.critic_mlp = nn.Sequential(
605 |             nn.Linear(output_dim*3, hidden_dim),
606 |             nn.ReLU(),
607 |             nn.Linear(hidden_dim, hidden_dim),
608 |             nn.ReLU(),
609 |             nn.Linear(hidden_dim, 1),
610 |             nn.LeakyReLU(alpha)
611 |         )
612 | 
613 |     def forward(self, u2s_path_loss, user_tasksize, server_comp_resource, edge_index):
614 |         
615 |         '''
616 |             input: path_losses of users to servers
617 |             output: power_sche, task_sche, comp_sche
618 |         '''
619 | 
620 |         feats = torch.cat([u2s_path_loss, user_tasksize, server_comp_resource], dim=0)
621 |         embeddings = self.encoder(feats)
622 | 
623 |         sches = self.sche_mlp(embeddings)
624 |         task_sche = sches[:self.output_dim]
625 |         task_sche = softmax(task_sche, index=edge_index[0])
626 | 
627 |         power_sche = sches[self.output_dim:2*self.output_dim]
628 |         power_sche = softmax(power_sche, edge_index[0])
629 | 
630 |         comp_sche = sches[2*self.output_dim:]
631 |         comp_sche = softmax(comp_sche, index=edge_index[1])
632 | 
633 |         sches = torch.cat([task_sche, power_sche, comp_sche], dim=-1)
634 |         critic_value = self.critic_mlp(sches.detach())
635 | 
636 |         return task_sche, power_sche, comp_sche, sches, critic_value
637 | 
638 | 
639 | 
640 | 
641 | 
642 | class MMSE(nn.Module):
643 |     def __init__(self, input_shape, args) -> None:
644 |         super(MMSE, self).__init__()
645 |         
646 |         self.args = args
647 | 
648 |         self.task_allocation = nn.Parameter(torch.rand(input_shape, requires_grad=True))
649 |         self.power_allocation = nn.Parameter(torch.rand(input_shape, requires_grad=True))
650 |         self.comp_allocation = nn.Parameter(torch.rand(input_shape, requires_grad=True))
651 | 
652 | 
653 |     def forward(self, compute_resource, path_losses, task_size, edge_index):
654 |         epsilon = 1e-9
655 |         extre = 1e-20
656 |         user_index = edge_index[0]      # s2u中源节点的索引
657 |         server_index = edge_index[1]    # s2u中目标节点的索引
658 | 
659 |         task_sche = softmax(self.task_allocation, edge_index[0])
660 |         power_sche = softmax(self.power_allocation, edge_index[0])
661 |         comp_sche = softmax(self.comp_allocation, edge_index[0])
662 | 
663 |         task_size = task_size[user_index]
664 |         # task_size = task_size[user_index]*args.tasksize_cof       # 重复采样映射到边中
665 |          
666 |         tasks = task_size * task_sche.squeeze()
667 |         compute_resource = compute_resource[server_index]
668 |         comp = compute_resource * comp_sche.squeeze()
669 | 
670 |         # power_sche = torch.clamp(power_sche, 1e-5, 1)
671 |         pw = power_sche.squeeze() * path_losses.squeeze()
672 | 
673 | 
674 |         pws_list = []
675 |         for idx in range(server_index[-1]+1):
676 |             # 取出所有到同一个server的pw加和
677 |             pws_idx = pw[torch.where(server_index==idx)].sum()
678 |             pws_list.append(pws_idx)
679 |         pws_list = torch.stack(pws_list)
680 |         pws_list = pws_list[server_index]
681 |         interference = pws_list-pw
682 |         rate = torch.log2(1 + torch.div(pw, interference+epsilon))
683 |         # rate = args.band_width * torch.log2(1+torch.div(pw, interference+epsilon))
684 |         
685 |         offloading_time = torch.div(tasks, rate+extre)
686 |         # offloading_time = torch.div(tasks, rate+extre) * (self.args.tasksize_cof/self.args.band_width)
687 | 
688 |         # compute_time = torch.div(tasks, comp+extre) * (self.args.tasksize_cof*self.args.cons_factor/self.args.comp_cof)
689 |         compute_time = torch.div(tasks, comp+extre)
690 | 
691 |         time_loss = offloading_time + compute_time
692 |         assert torch.isnan(time_loss).sum()==0
693 | 
694 |         time_loss_list = []
695 |         for idx in range(user_index[-1]+1):
696 |             tl_idx = time_loss[torch.where(user_index==idx)].sum()
697 |             time_loss_list.append(tl_idx)
698 |         time_loss_list = torch.stack(time_loss_list)
699 | 
700 |         return time_loss_list.mean()
701 | 


--------------------------------------------------------------------------------
/pretrain.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | from arguments import args
  3 | import numpy as np
  4 | from torch_geometric.loader import DataLoader
  5 | from layouts import generate_layouts
  6 | from off_loading_models import TaskLoad, PCNet
  7 | from tqdm import tqdm
  8 | import random
  9 | 
 10 | def compute_loss_nn(task_allocation, power_allocation, comp_allocation, task_size, compute_resource, path_losses, user_index, server_index):
 11 |     
 12 |     # task_size : vector N
 13 |     # task_allocation: mat pop_num x 3*M*N
 14 |     # index: vector 3*M*N
 15 |     
 16 |     epsilon = 1e-9
 17 |     extre = 1e-20
 18 |     server_index_first = server_index.reshape((batch_size, -1))[0]
 19 |     user_index_first = user_index.reshape((batch_size, -1))[0]
 20 |     # user_index = edge_index[0]      # s2u中源节点的索引
 21 |     # server_index = edge_index[1]    # s2u中目标节点的索引
 22 |     
 23 |     # power_allocation = torch.clamp(power_allocation, 1e-5, 1)
 24 |     pw_ini = power_allocation * path_losses    # mat pop_num x M*N
 25 |     # pw = torch.clamp(pw, 1e-5, 1)
 26 |     
 27 |     # 将信道状态过小的设置为0
 28 |     mask_pw = torch.where(pw_ini<args.pw_threshold)
 29 |     
 30 |     pw = pw_ini.clone()
 31 | 
 32 |     pw[mask_pw] = 0
 33 | 
 34 |     comp_allocation_clone = comp_allocation.clone()
 35 |     comp_allocation_clone[mask_pw] = 0
 36 |     comp_allocation_normed = torch.zeros((comp_allocation_clone.shape[0], comp_allocation_clone.shape[1], server_index_first[-1]+1), device=args.device)
 37 |     comp_allocation_normed.scatter_(2, server_index_first.repeat((batch_size, 1)).unsqueeze(2), comp_allocation_clone.unsqueeze(2))
 38 |     comp_allocation_normed = comp_allocation_normed.sum(1)[:, server_index_first]
 39 |     comp_allocation_normed = torch.div(comp_allocation_clone, comp_allocation_normed+extre)
 40 | 
 41 |     task_allocation_clone = task_allocation.clone()
 42 |     task_allocation_clone[mask_pw] = 0
 43 |     task_allocation_normed = torch.zeros((task_allocation_clone.shape[0], task_allocation_clone.shape[1], user_index_first[-1]+1), device=args.device)
 44 |     task_allocation_normed.scatter_(2, user_index_first.repeat((batch_size, 1)).unsqueeze(2), task_allocation_clone.unsqueeze(2))
 45 |     task_allocation_normed = task_allocation_normed.sum(1)[:, user_index_first]
 46 |     task_allocation_normed = torch.div(task_allocation_clone, task_allocation_normed+extre)
 47 | 
 48 |     # 计算速率
 49 |     pw_list = torch.zeros((pw.shape[0], pw.shape[1], server_index_first[-1]+1), device=args.device)   # mat pop_num x MN x N
 50 |     pw_list.scatter_(2, server_index_first.repeat((batch_size, 1)).unsqueeze(2), pw.unsqueeze(2))
 51 |     pws_list = pw_list.sum(1)[:, server_index_first]  # mat pop_num x MN
 52 | 
 53 |     interference = pws_list-pw
 54 |     rate = torch.log2(1+torch.div(pw, interference+epsilon))
 55 |     # rate = args.band_width * torch.log2(1+torch.div(pw, interference+epsilon))
 56 | 
 57 | 
 58 |     task_size = task_size[:, user_index_first]       # M*N    
 59 |     # task_size = task_size[user_index]*args.tasksize_cof       # 重复采样映射到边中
 60 |     tasks = task_size * task_allocation_normed   # mat pop_num x M*N
 61 | 
 62 |     compute_resource = compute_resource[:, server_index_first]
 63 |     # compute_resource = compute_resource[server_index]*args.comp_cof       # 
 64 | 
 65 |     comp = compute_resource * comp_allocation_normed
 66 | 
 67 | 
 68 |     # offloading_time = torch.div(tasks, rate+extre) * (args.tasksize_cof/args.band_width)
 69 |     offloading_time = torch.div(tasks, rate+extre)
 70 | 
 71 |     # compute_time = torch.div(tasks, comp+extre) * (args.tasksize_cof*args.cons_factor/args.comp_cof)
 72 |     compute_time = torch.div(tasks, comp+extre)
 73 | 
 74 |     time_loss = offloading_time + compute_time      # pop_num x MN
 75 |     assert torch.isnan(time_loss).sum()==0
 76 | 
 77 | 
 78 |     time_loss_list = torch.zeros((time_loss.shape[0], time_loss.shape[1], user_index_first[-1]+1), device=args.device)
 79 |     time_loss_list.scatter_(2, user_index_first.repeat((batch_size, 1)).unsqueeze(2), time_loss.unsqueeze(2))
 80 |     time_loss_list = time_loss_list.sum(1)      # pop_num x MN
 81 | 
 82 |     return time_loss_list.mean()
 83 | 
 84 | 
 85 | def compute_loss(task_allocation, power_allocation, comp_allocation, compute_resource, path_losses, task_size, user_index, server_index):
 86 |     
 87 |     epsilon = 1e-9
 88 |     extre = 1e-20
 89 |     # user_index = edge_index[0]      # s2u中源节点的索引
 90 |     # server_index = edge_index[1]    # s2u中目标节点的索引
 91 |     
 92 |     pw_ini = power_allocation.squeeze() * path_losses.squeeze()
 93 |     # pw小于阈值的对应power设为0
 94 |     pw = pw_ini.clone()
 95 |     # mask_pw = torch.where(pw_ini<args.pw_threshold)
 96 |     # pw[mask_pw] = 0
 97 | 
 98 |     pw_user_list = torch.zeros((len(pw), user_index[-1]+1), device=args.device)
 99 |     # pw_user_ini_list = torch.zeros((len(pw), user_index[-1]+1), device=args.device)
100 |     # pw_user_ini_list.scatter_(1, user_index.unsqueeze(1), pw_ini.unsqueeze(1))
101 |     pw_user_list.scatter_(1, user_index.unsqueeze(1), pw_ini.unsqueeze(1))
102 |     # 如果某一个user的发射功率均位于阈值以下
103 |     pw_masked = pw_user_list.clone()
104 |     pw_masked[torch.where(pw_masked < args.pw_threshold)] = 0
105 |     invalid_index = torch.where(pw_masked.sum(0)==0)[0]   # 是否有对所有server都低于阈值的
106 |     # assert len(invalid_index)==0
107 |     max_pw_index = pw_user_list[:, invalid_index].argmax(0) # 对所有server的pw都低于阈值的user 取信号最强的server
108 |     pw_masked[max_pw_index, invalid_index] = pw_user_list[max_pw_index, invalid_index]
109 |     pw = pw_masked.sum(1)
110 |     mask_pw = torch.where(pw==0)
111 | 
112 |     pw_list = torch.zeros((len(pw), server_index[-1]+1), device=args.device)
113 |     pw_list.scatter_(1, server_index.unsqueeze(1), pw.unsqueeze(1))
114 |     
115 |     pws_list = pw_list.sum(0)[server_index]
116 | 
117 | 
118 |     interference = pws_list-pw
119 |     rate = torch.log2(1+torch.div(pw, interference+epsilon))
120 | 
121 |     task_allocation_clone = task_allocation.clone().squeeze() + 1e-8
122 |     task_allocation_clone[mask_pw] = 0
123 |     task_allocation_normed = torch.zeros((len(task_allocation_clone), user_index[-1]+1), device=args.device)
124 |     task_allocation_normed.scatter_(1, user_index.unsqueeze(1), task_allocation_clone.unsqueeze(1))
125 |     # assert len(torch.where(task_allocation_normed.sum(0)==0)[0]) == 0
126 |     task_allocation_normed_2 = task_allocation_normed.sum(0)[user_index]
127 |     task_allocation_final = torch.div(task_allocation_clone, task_allocation_normed_2+extre)
128 |     # task_allocation_clone =  softmax(task_allocation_clone, user_index)
129 | 
130 |     task_size = task_size[user_index]
131 |     # task_size = task_size[user_index]*args.tasksize_cof       # 重复采样映射到边中
132 |     
133 |     tasks = task_size * task_allocation_final
134 | 
135 |     comp_allocation_clone = comp_allocation.clone().squeeze()
136 |     comp_allocation_clone[mask_pw] = 0
137 |     comp_allocation_normed = torch.zeros((len(comp_allocation_clone), server_index[-1]+1), device=args.device)
138 |     comp_allocation_normed.scatter_(1, server_index.unsqueeze(1), comp_allocation_clone.unsqueeze(1))
139 |     comp_allocation_normed_2 = comp_allocation_normed.sum(0)[server_index]
140 |     comp_allocation_final = torch.div(comp_allocation_clone, comp_allocation_normed_2+extre)
141 |     # compute_resource = compute_resource[server_index]*args.comp_cof       # 
142 | 
143 |     compute_resource = compute_resource[server_index]
144 |     comp = compute_resource * comp_allocation_final
145 | 
146 | 
147 |     # offloading_time = torch.div(tasks, rate+extre) * (args.tasksize_cof/args.band_width)
148 |     offloading_time = torch.div(tasks, rate+extre)
149 | 
150 |     # compute_time = torch.div(tasks, comp+extre) * (args.tasksize_cof*args.cons_factor/args.comp_cof)
151 |     compute_time = torch.div(tasks, comp+extre)
152 | 
153 |     # compute_time = torch.clamp(compute_time, -1, 3000)
154 |     # offloading_time = torch.clamp(offloading_time, -1, 3000)
155 | 
156 |     time_loss = offloading_time + compute_time
157 |     assert torch.isnan(time_loss).sum()==0
158 | 
159 |     time_loss_list = torch.zeros((len(time_loss), user_index[-1]+1), device=args.device)
160 |     time_loss_list.scatter_(1, user_index.unsqueeze(1), time_loss.unsqueeze(1))
161 |     time_loss_list = time_loss_list.sum(0)
162 | 
163 |     return time_loss_list.mean()
164 | 
165 | def HGNN_train(model, train_loader, test_loader):
166 |     policy_losses = []
167 |     test_policy_losses = []
168 |     optimizer = torch.optim.Adam(model.parameters(), lr=hgnn_lr)
169 |     optimizer_stepLR = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20,  gamma=0.9)
170 |     for time_step in tqdm(range(Epochs)):
171 |         # training
172 |         
173 |         loss_sum = 0
174 |         length = 0
175 |         for graph in train_loader:    # graph为一个batch
176 |             task_allocation, power_allocation, comp_allocation = model(graph.x_dict, graph.edge_index_dict, graph.edge_attr_dict)
177 |             user_index = graph['user', 'u2s', 'server'].edge_index[0]
178 |             server_index = graph['user', 'u2s', 'server'].edge_index[1]
179 |             loss_batch = compute_loss(task_allocation, power_allocation, comp_allocation, graph['server'].x[:, 0], graph['user', 'u2s', 'server'].path_loss, graph['user'].x[:, 0], user_index, server_index)
180 |             optimizer.zero_grad()
181 |             loss_batch.backward()
182 |             optimizer.step()
183 |             loss_sum += loss_batch.item()
184 |             length += 1
185 |         optimizer_stepLR.step()
186 |         policy_loss = loss_sum/length
187 |         if (time_step+1) % eval_fre == 0:
188 |             eval_loss = HGNN_eval(time_step+1, model, test_loader)
189 |             test_policy_losses.append(eval_loss)
190 |             model.train()
191 |         if (time_step + 1) % save_fre == 0:
192 |             torch.save(model, './TO_models/Pretrain/'+multi_scales+'/HGNN_{}_{}.pt'.format(train_num_layouts, time_step+1))
193 |         policy_losses.append(policy_loss)
194 |         print('step=={}, policy_loss=={}'.format(time_step, policy_loss))
195 |     return model, np.array(policy_losses), np.array(test_policy_losses)
196 | 
197 | def HGNN_eval(time_step, model, loader):
198 |     loss_sum = 0
199 |     length = 0        
200 |     model.eval()
201 |     for graph in loader:    # graph为一个batch
202 |         task_allocation, power_allocation, comp_allocation = model(graph.x_dict, graph.edge_index_dict, graph.edge_attr_dict)
203 |         user_index = graph['user', 'u2s', 'server'].edge_index[0]
204 |         server_index = graph['user', 'u2s', 'server'].edge_index[1]
205 |         loss_batch = compute_loss(task_allocation, power_allocation, comp_allocation, graph['server'].x[:, 0], graph['user', 'u2s', 'server'].path_loss, graph['user'].x[:, 0], user_index, server_index)
206 |     
207 |         loss_sum += loss_batch.item()
208 |         length += 1
209 |     policy_loss = loss_sum/length
210 | 
211 |     print('step=={}, evaluate_policy_loss=={}'.format(time_step, policy_loss))
212 |     return policy_loss
213 | 
214 | 
215 | def NN_train(train_loader, test_loader):
216 |     server_num = min_server
217 |     sum_loss = 0
218 |     length = 0
219 |     train_losses = []
220 |     test_losses = []
221 | 
222 |     model = PCNet((server_num**2)*3+server_num*4, args.hidden_dim, (server_num**2)*3, args.alpha).to(args.device)
223 |     optimizer = torch.optim.Adam(model.parameters(), lr=pcnet_lr)
224 |     schedular = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.9)
225 |     for step in tqdm(range(Epochs)):
226 |         for idx, batch in enumerate(train_loader):
227 |             u2s_index = batch['user', 'u2s', 'server'].edge_index
228 |             user_index = u2s_index[0]
229 |             server_index = u2s_index[1]
230 |             u2s_path_loss = batch['user', 'u2s', 'server'].path_loss.squeeze().reshape((batch_size, -1))
231 |             u2s_path_loss_feat = batch['user', 'u2s', 'server'].edge_attr.squeeze().reshape((batch_size, -1))
232 |             user_tasksize = batch['user'].x[:, 0].reshape((batch_size, -1))
233 |             server_comp_resource = batch['server'].x[:, 0].reshape((batch_size, -1))
234 | 
235 |             task_sche, power_sche, comp_sche = model(u2s_path_loss_feat, user_tasksize, server_comp_resource, u2s_index)
236 |             loss = compute_loss_nn(task_sche, power_sche, comp_sche, user_tasksize, server_comp_resource, u2s_path_loss, user_index, server_index)
237 |             loss.backward()
238 |             optimizer.step()
239 |             optimizer.zero_grad()
240 |             sum_loss += loss.item()
241 |             length += 1
242 |         # schedular.step()
243 |         mean_loss = sum_loss/length
244 |         train_losses.append(mean_loss)
245 |         if (step+1) % eval_fre == 0:
246 |             eval_loss = NN_eval(step+1, test_loader, model)
247 |             test_losses.append(eval_loss)
248 |             model.train()
249 | 
250 |         if (step + 1) % save_fre == 0:
251 |             torch.save(model, './TO_models/Pretrain/'+multi_scales+'/PcNet_{}_{}.pt'.format(train_num_layouts, step+1))
252 | 
253 |         print('Time step: {} \t\t PCNet training loss: {}'.format(step, mean_loss))
254 |         
255 | 
256 |     return model, np.array(train_losses), np.array(test_losses)
257 | 
258 | def NN_eval(step, loader, model):
259 |     sum_loss = 0
260 |     length = 0
261 |     model.eval()
262 |     for idx, batch in enumerate(loader):
263 |         u2s_index = batch['user', 'u2s', 'server'].edge_index
264 |         user_index = u2s_index[0]
265 |         server_index = u2s_index[1]
266 |         u2s_path_loss = batch['user', 'u2s', 'server'].path_loss.squeeze().reshape((batch_size, -1))
267 |         u2s_path_loss_feat = batch['user', 'u2s', 'server'].edge_attr.squeeze().reshape((batch_size, -1))
268 |         user_tasksize = batch['user'].x[:, 0].reshape((batch_size, -1))
269 |         server_comp_resource = batch['server'].x[:, 0].reshape((batch_size, -1))
270 | 
271 |         task_sche, power_sche, comp_sche = model(u2s_path_loss_feat, user_tasksize, server_comp_resource, u2s_index)
272 |         loss = compute_loss_nn(task_sche, power_sche, comp_sche, user_tasksize, server_comp_resource, u2s_path_loss,  user_index, server_index)
273 |         sum_loss += loss.item()
274 |         length += 1
275 |     mean_loss = sum_loss/length
276 |     print('Time step: {}, PCNet evaluate loss: {}'.format(step, mean_loss))
277 |     return mean_loss
278 | 
279 | 
280 | def setup_seed(seed):
281 |      torch.manual_seed(seed)
282 |      torch.cuda.manual_seed_all(seed)
283 |      np.random.seed(seed)
284 |      random.seed(seed)
285 |      torch.backends.cudnn.deterministic = True
286 |      np.random.seed(seed)
287 | 
288 | 
289 | 
290 | if __name__ == '__main__':
291 |     setup_seed(20)
292 |     Epochs =500
293 |     hgnn_lr = 5e-4
294 |     pcnet_lr = 1e-3
295 |     train_num_layouts = 2048
296 |     test_num_layouts = 512
297 |     batch_size = 32
298 |     eval_fre = 5
299 |     save_fre = 50
300 |     multi_scales = 'large'
301 |     if multi_scales == 'small':
302 |         min_server = 2
303 |         max_server = 10
304 |     elif multi_scales == 'medium':
305 |         min_server = 15
306 |         max_server = 25
307 |     elif multi_scales == 'large':
308 |         min_server = 25
309 |         max_server = 35
310 |     train_server_nums = np.random.randint(min_server, max_server+1, train_num_layouts)
311 |     train_user_nums = 3*train_server_nums
312 |     test_server_nums = np.random.randint(min_server, max_server+1, test_num_layouts)
313 |     test_user_nums = 3*test_server_nums
314 | 
315 |     nn_train_server_nums = np.random.randint(min_server, min_server+1, train_num_layouts)
316 |     nn_train_user_nums = 3*nn_train_server_nums
317 |     nn_test_server_nums = np.random.randint(min_server, min_server+1, train_num_layouts)
318 |     nn_test_user_nums = 3*nn_test_server_nums
319 | 
320 | 
321 |     # env_max_length = np.sqrt(10*300)
322 |     gnn_train_layouts = generate_layouts(train_user_nums, train_server_nums, args)
323 |     gnn_test_layouts = generate_layouts(test_user_nums, test_server_nums, args)
324 |     nn_train_layouts = generate_layouts(nn_train_user_nums, nn_train_server_nums, args)
325 |     nn_test_layouts = generate_layouts(nn_test_user_nums, nn_test_server_nums, args)
326 |     
327 |     
328 |     gnn_train_loader = DataLoader(gnn_train_layouts, batch_size=batch_size, shuffle=True)
329 |     gnn_test_loader = DataLoader(gnn_test_layouts, batch_size=batch_size, shuffle=True)
330 | 
331 |     nn_train_loader = DataLoader(nn_train_layouts, batch_size=batch_size, shuffle=True)
332 |     nn_test_loader = DataLoader(nn_test_layouts, batch_size=batch_size, shuffle=True)
333 | 
334 |     # task offloading trainning process
335 |     Ol_Model = TaskLoad(args.num_layers, args.input_dim, args.hidden_dim, args.max_server_num, args.alpha).to(args.device)
336 |     hgnn_model, train_loss, test_loss = HGNN_train(Ol_Model, gnn_train_loader, gnn_test_loader)
337 | 
338 |     pcnet_model, nn_train_losses, nn_test_losses = NN_train(nn_train_loader, nn_test_loader)
339 | 


--------------------------------------------------------------------------------