修改生成数据方式和网络训练方式

base_optimizer/optimizer_common.py

@@ -173,7 +173,7 @@ def timer_wrapper(func):
         start_time = time.time()
         result = func(*args, **kwargs)
 
-        print("function {} running time :  {} s".format(func.__name__, time.time() - start_time))
+        print(f"function {func.__name__} running time :  {time.time() - start_time:.3f} s")
         return result
 
     return measure_time

base_optimizer/optimizer_feederpriority.py

@@ -66,7 +66,6 @@ def feeder_allocate(component_data, pcb_data, feeder_data, figure=False):
     nozzle_assigned_counter = optimal_nozzle_assignment(component_data, pcb_data)
     head_assign_indexes = list(range(max_head_index))
     nozzle_pattern, optimal_nozzle_pattern, optimal_nozzle_points = [], None, 0
-    # nozzle_pattern = ['CN220', 'CN065','CN065','CN065','CN065','CN220']
 
     # 先排序
     nozzle_pattern_list = []

base_optimizer/optimizer_hybridgenetic.py

@@ -294,8 +294,6 @@ def convert_individual_2_result(component_data, component_point_pos, designated_
 
 @timer_wrapper
 def optimizer_hybrid_genetic(pcb_data, component_data, hinter=True):
-    random.seed(0)
-    np.random.seed(0)
     nozzle_assigned_counter = optimal_nozzle_assignment(component_data, pcb_data)
 
     # nozzle assignment result:

base_optimizer/optimizer_interface.py

@@ -80,7 +80,7 @@ def base_optimizer(machine_index, pcb_data, component_data, feeder_data=None, me
 
     if hinter:
         optimization_assign_result(component_data, pcb_data, component_result, cycle_result, feeder_slot_result,
-                                   nozzle_hinter=True, component_hinter=False, feeder_hinter=True)
+                                   nozzle_hinter=True, component_hinter=True, feeder_hinter=True)
 
         print('----- Placement machine ' + str(machine_index) + ' ----- ')
         print('-Cycle counter: {}'.format(info.cycle_counter))

generator.py

@@ -1,3 +1,6 @@
+import random
+
+import numpy as np
 import pandas as pd
 
 from base_optimizer.optimizer_common import *
@@ -8,53 +11,58 @@ class DataMgr:
         self.min_placement_points = 100
         self.max_placement_points = 800
 
-        self.max_component_types = 50
+        self.max_component_types = 30
         self.default_feeder_limit = 1
         self.nozzle_type_list = ['CN065', 'CN140', 'CN220', 'CN040']
 
-        self.x_range = [50, 100, 150, 200, 300, 400, 500]
+        # self.x_range = [50, 100, 150, 200, 300, 400, 500]
-        self.y_range = [50, 100, 150, 200, 300, 400, 500]
+        # self.y_range = [50, 100, 150, 200, 300, 400, 500]
+
+        self.x_range = [400]
+        self.y_range = [200]
 
         self.counter = 0
+        self.update = 10
         self.pre_file = None
+        self.part_col = ["part", "desc", "fdr", "nz", 'camera', 'group', 'feeder-limit', 'points']
+        self.component_data = pd.DataFrame(columns=self.part_col)  # the component list update for several rounds
 
-    def generator(self):
+    def generator(self, mode='Train'):
         boundary = [random.choice(self.x_range), random.choice(self.y_range)]
-        total_points = random.randint(self.min_placement_points, self.max_placement_points)
 
         # determine the nozzle type of component
-        component_list = defaultdict(str)
+        if self.counter % 10 == 0 or mode == 'test':
-        for cp_idx in range(min(random.randint(1, self.max_component_types), total_points)):
+            self.component_data = self.component_data.loc[[]]
-            component_list['C' + str(cp_idx)] = random.choice(self.nozzle_type_list)
+            total_points = random.randint(self.min_placement_points, self.max_placement_points)
+            total_nozzles = random.randint(1, len(self.nozzle_type_list))
+            selected_nozzle = random.sample(self.nozzle_type_list, total_nozzles)
+            for cp_idx in range(min(random.randint(1, self.max_component_types), total_points)):
+                part, nozzle = 'C' + str(cp_idx), random.choice(selected_nozzle)
+                self.component_data = pd.concat([self.component_data, pd.DataFrame(
+                    [part, '', 'SM8', nozzle, '飞行相机1', 'CHIP-Rect', self.default_feeder_limit, 0],
+                    index=self.part_col).T], ignore_index=True)
+
+            random_fractions = np.random.rand(len(self.component_data))
+            normalized_fractions = random_fractions / random_fractions.sum()
+            for cp_idx, fraction in enumerate(normalized_fractions):
+                self.component_data.iloc[cp_idx].points = round(fraction * total_points)
 
         step_col = ["ref", "x", "y", "z", "r", "part", "desc", "fdr", "nz", "hd", "cs", "cy", "sk", "bl", "ar", "pl", "lv"]
         pcb_data = pd.DataFrame(columns=step_col)
+        idx = 1
+        for _, data in self.component_data.iterrows():
+            for _ in range(data.points):
+                part, nozzle = data.part, data.nz
+                pos_x, pos_y = np.random.uniform(0, boundary[0]), np.random.uniform(0, boundary[1])
+                pcb_data = pd.concat([pcb_data, pd.DataFrame([['R' + str(idx), -pos_x, pos_y,
+                                                        0.000, 0.000, part, '', 'A', '1-0 ' + nozzle, 1, 1, 1, 0,
+                                                        1, 1, 1, 'L0']], columns=pcb_data.columns)], ignore_index=True)
+                idx += 1
 
-        for idx in range(total_points):
-            part = random.choice(list(component_list.keys()))
-            nozzle = component_list[part]
-
-            pos_x, pos_y = np.random.uniform(0, boundary[0]), np.random.uniform(0, boundary[1])
-            pcb_data = pd.concat([pcb_data, pd.DataFrame([['R' + str(idx), -pos_x, pos_y,
-                                                       0.000, 0.000, part, '', 'A', '1-0 ' + nozzle, 1, 1, 1, 0,
-                                                       1, 1, 1, 'L0']], columns=pcb_data.columns)], ignore_index=True)
-
-        part_col = ["part", "desc", "fdr", "nz", 'camera', 'group', 'feeder-limit', 'points']
-        component_data = pd.DataFrame(columns=part_col)
-
-        for _, data in pcb_data.iterrows():
-            part, nozzle = data.part, data.nz.split(' ')[1]
-            if part not in component_data['part'].values:
-                component_data = pd.concat([component_data, pd.DataFrame(
-                    [part, '', 'SM8', nozzle, '飞行相机1', 'CHIP-Rect', self.default_feeder_limit, 0], index=part_col).T],
-                                           ignore_index=True)
-
-            part_index = component_data[component_data['part'] == part].index.tolist()[0]
-            component_data.loc[part_index, 'points'] += 1
         self.counter += 1
-        return pcb_data, component_data
+        return pcb_data, self.component_data
 
-    def recorder(self, file_path, info: OptInfo, pcb_data, component_data):
+    def recorder(self, file_handle, info: OptInfo, pcb_data, component_data):
         lineinfo = '{:.6f}'.format(info.placement_time) + '\t' + str(info.cycle_counter) + '\t' + str(
             info.nozzle_change_counter) + '\t' + str(info.pickup_counter) + '\t' + '{:.3f}'.format(
             info.pickup_movement) + '\t' + '{:.3f}'.format(info.placement_movement)
@@ -62,6 +70,11 @@ class DataMgr:
         lineinfo += '\t' + '{:.3f}'.format(pcb_data['x'].max() - pcb_data['x'].min()) + '\t' + '{:.3f}'.format(
             pcb_data['y'].max() - pcb_data['y'].min())
 
+        part_xposition, part_yposition = defaultdict(list), defaultdict(list)
+        for _, data in pcb_data.iterrows():
+            part_xposition[data['part']].append(data['x'])
+            part_yposition[data['part']].append(data['y'])
+
         point_counter, component_counter = 0, 0
         nozzle_type = set()
         for _, data in component_data.iterrows():
@@ -75,11 +88,12 @@ class DataMgr:
 
         for _, data in component_data.iterrows():
             lineinfo += '\t' + data.part + '\t' + data.nz + '\t' + str(data.points)
-        lineinfo += '\n'
+            # lineinfo += '\t' + str(
+            #     round((np.average(part_xposition[data.part]) + stopper_pos[0] - slotf1_pos[0]) / slot_interval))
+
+        lineinfo += '\n'
+        file_handle.write(lineinfo)
 
-        with open(file_path, 'a') as f:
-            f.write(lineinfo)
-        f.close()
 
     def saver(self, file_path: str, pcb_data):
         lineinfo = ''
@@ -103,12 +117,23 @@ class DataMgr:
         for idx, nozzle in enumerate(self.nozzle_type_list):
             cp2nz[nozzle] = idx
 
+        # === general info ===
         total_points = sum(points for points in cp_points.values())
         total_component_types, total_nozzle_types = len(cp_points.keys()), len(set(cp_nozzle.values()))
         data = [total_points, total_component_types, total_nozzle_types]
         data.extend([width, height])
 
+        # === nozzle info ===
+        data_slice = [0 for _ in range(len(self.nozzle_type_list))]
         for component, points in cp_points.items():
+            idx = cp2nz[cp_nozzle[component]]
+            data_slice[idx] += points
+        data.extend(data_slice)
+
+        # === component info ===
+        cp_items = [[component, points] for component, points in cp_points.items()]
+        cp_items = sorted(cp_items, key=lambda x: (-x[1], x[0]))
+        for component, points in cp_items:
             nozzle = cp_nozzle[component]
 
             data_slice = [0 for _ in range(len(self.nozzle_type_list))]
@@ -120,15 +145,43 @@ class DataMgr:
 
         return data
 
+    def decode(self, line_info):
+        boundary = [random.choice(self.x_range), random.choice(self.y_range)]
+        items = line_info.split('\t')
+        total_points, total_component_types = int(items[8]), int(items[9])
+
+        part_col = ["part", "desc", "fdr", "nz", 'camera', 'group', 'feeder-limit', 'points']
+        step_col = ["ref", "x", "y", "z", "r", "part", "desc", "fdr", "nz", "hd", "cs", "cy", "sk", "bl", "ar", "pl",
+                    "lv"]
+
+        component_data = pd.DataFrame(columns=part_col)
+        pcb_data = pd.DataFrame(columns=step_col)
+
+        idx = 1
+        for cp_counter in range(total_component_types):
+            # todo: 这里为了调试暂时未修改
+            part, nozzle = items[11 + cp_counter * 3], items[12 + cp_counter * 3]
+            points = int(items[13 + cp_counter * 3])
+
+            component_data = pd.concat([component_data, pd.DataFrame(
+                [part, '', 'SM8', nozzle, '飞行相机1', 'CHIP-Rect', self.default_feeder_limit, points], index=part_col).T],
+                                       ignore_index=True)
+
+            for _ in range(points):
+                pos_x, pos_y = np.random.uniform(0, boundary[0]), np.random.uniform(0, boundary[1])
+                pcb_data = pd.concat([pcb_data, pd.DataFrame([['R' + str(idx), -pos_x, pos_y, 0.000, 0.000, part, '',
+                                                               'A', '1-0 ' + nozzle, 1, 1, 1, 0, 1, 1, 1, 'L0']],
+                                                             columns=pcb_data.columns)], ignore_index=True)
+        return pcb_data, component_data
+
     def loader(self, file_path):
         train_data, time_data = [], []
         cycle_data, nozzle_change_data, pickup_data, movement_data, point_data = [], [], [], [], []
-        pcb_width, pcb_height = [], []
         with open(file_path, 'r') as file:
             line = file.readline()
             while line:
                 items = line.split('\t')
-                total_points, total_component_types = int(items[8]), int(items[9])
+                total_points, total_component_types = float(items[8]), float(items[9])
 
                 cycle_data.append(float(items[1]))
                 nozzle_change_data.append(float(items[2]))
@@ -140,7 +193,7 @@ class DataMgr:
                 time_data.append(float(items[0]))
 
                 cp_points, cp_nozzle = defaultdict(int), defaultdict(str)
-                for cp_counter in range(total_component_types):
+                for cp_counter in range(int(total_component_types)):
                     component_type, nozzle_type = items[11 + cp_counter * 3], items[12 + cp_counter * 3]
                     points = int(items[13 + cp_counter * 3])
 
@@ -152,5 +205,8 @@ class DataMgr:
         return train_data, time_data, cycle_data, nozzle_change_data, pickup_data, movement_data, point_data
 
     def get_feature(self):
-        return self.max_component_types * len(self.nozzle_type_list) + 5
+        return (self.max_component_types + 1) * len(self.nozzle_type_list) + 5
+
+    def get_update_round(self):
+        return self.update
 

optimizer.py

@@ -24,7 +24,7 @@ def optimizer(pcb_data, component_data, line_optimizer, machine_optimizer, machi
         return
 
     assignment_result_cpy = copy.deepcopy(assignment_result)
-    placement_points, placement_time = [], []
+    placement_points, assembly_info = [], []
     partial_pcb_data, partial_component_data = defaultdict(pd.DataFrame), defaultdict(pd.DataFrame)
     for machine_index in range(machine_number):
         partial_pcb_data[machine_index] = pd.DataFrame(columns=pcb_data.columns)
@@ -128,24 +128,33 @@ def optimizer(pcb_data, component_data, line_optimizer, machine_optimizer, machi
         for info in part_info:
             partial_component_data[machine_index].loc[info[0], 'feeder-limit'] = math.ceil(info[2] / max_avl_feeder)
 
-        placement_time.append(base_optimizer(machine_index + 1, data, partial_component_data[machine_index],
+        assembly_info.append(base_optimizer(machine_index + 1, data, partial_component_data[machine_index],
                                              feeder_data=pd.DataFrame(columns=['slot', 'part', 'arg']),
-                                             method=machine_optimizer, hinter=True).placement_time)
+                                             method=machine_optimizer, hinter=True))
 
-    average_time, standard_deviation_time = sum(placement_time) / machine_number, 0
+    with open('model/lr_model.pkl', 'rb') as f:
+        lr = pickle.load(f)
+
+    average_time, standard_deviation_time = sum(
+        [assembly_info[m].placement_time for m in range(machine_number)]) / machine_number, 0
     for machine_index in range(machine_number):
-        total_component_types = 0
+        total_component_types = sum(1 if pt else 0 for pt in assignment_result_cpy[machine_index])
-        for points in assignment_result_cpy[machine_index]:
+        placement_time = assembly_info[machine_index].placement_time
-            if points:
+
-                total_component_types += 1
+        regression_time = lr.coef_[0][0] * assembly_info[machine_index].cycle_counter + lr.coef_[0][1] * assembly_info[
-        print('assembly time for machine ' + str(machine_index + 1) + ': ' + str(
+            machine_index].nozzle_change_counter + lr.coef_[0][2] * assembly_info[machine_index].pickup_counter + \
-            placement_time[machine_index]) + ' s, ' + 'total placements: ' + str(placement_points[machine_index])
+                          lr.coef_[0][3] * assembly_info[machine_index].pickup_movement + lr.coef_[0][4] * \
-              + ', total component types: ' + str(total_component_types))
+                          placement_points[machine_index] + lr.intercept_[0]
-        standard_deviation_time += pow(placement_time[machine_index] - average_time, 2)
+
+        print(f'assembly time for machine {machine_index + 1: d}: {placement_time: .3f} s, total placement: '
+              f'{placement_points[machine_index]}, total component types {total_component_types: d}', end=', ')
+        print(f'regression time: {regression_time: .3f} s')
+        standard_deviation_time += pow(placement_time - average_time, 2)
     standard_deviation_time /= machine_number
     standard_deviation_time = math.sqrt(standard_deviation_time)
 
-    print('finial assembly time: ' + str(max(placement_time)) + 's, standard deviation: ' + str(standard_deviation_time))
+    print(f'finial assembly time: {max(info.placement_time for info in assembly_info): .3f} s, '
+          f'standard deviation: {standard_deviation_time: .3f}')
 
 
 @timer_wrapper
@@ -169,6 +178,7 @@ def main():
     # 加载PCB数据
     pcb_data, component_data, _ = load_data(params.filename, default_feeder_limit=params.feeder_limit,
                                             cp_auto_register=params.auto_register, load_feeder_data=False)  # 加载PCB数据
+
     optimizer(pcb_data, component_data, params.line_optimizer, params.machine_optimizer, params.machine_number)
 
 

optimizer_genetic.py

@@ -283,7 +283,6 @@ def assemblyline_optimizer_genetic(pcb_data, component_data, machine_number):
     # the number of generation: 500
     crossover_rate, mutation_rate = 0.8, 0.1
     population_size, n_generations = 200, 500
-    # population_size, n_generations = 30, 50
 
     # the number of placement points, the number of available feeders, and nozzle type of component respectively
     component_points, component_feeders, component_nozzle = defaultdict(int), defaultdict(int), defaultdict(str)
@@ -300,7 +299,7 @@ def assemblyline_optimizer_genetic(pcb_data, component_data, machine_number):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     net = Net(input_size=data_mgr.get_feature(), output_size=1).to(device)
 
-    net.load_state_dict(torch.load('model_state.pth'))
+    net.load_state_dict(torch.load('model/net_model.pth'))
     # optimizer = torch.optim.Adam(net.parameters(), lr=0.1)
     # optimizer.load_state_dict(torch.load('optimizer_state.pth'))
 
@@ -361,6 +360,7 @@ def assemblyline_optimizer_genetic(pcb_data, component_data, machine_number):
     best_individual = population[np.argmax(pop_val)]
     val, assignment_result = cal_individual_val(component_points, component_feeders, component_nozzle, machine_number,
                                                 best_individual, data_mgr, net)
+
     print('final value: ', val)
     # available feeder check
     for part_index, data in component_data.iterrows():

optimizer_hyperheuristic.py

@@ -1,6 +1,8 @@
+import os
 import pickle
 
 import numpy as np
+import torch.nn
 
 from base_optimizer.optimizer_interface import *
 from generator import *
@@ -9,11 +11,11 @@ os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
 
 
 class Net(torch.nn.Module):
-    def __init__(self, input_size, output_size):
+    def __init__(self, input_size, hidden_size=1024, output_size=1):
         super(Net, self).__init__()
-        self.fc1 = torch.nn.Linear(input_size, 1024)
+        self.fc1 = torch.nn.Linear(input_size, hidden_size)
         self.relu = torch.nn.ReLU()  # 激活函数
-        self.fc2 = torch.nn.Linear(1024, output_size)
+        self.fc2 = torch.nn.Linear(hidden_size, output_size)
 
     def forward(self, x):
         x = self.fc1(x)
@@ -22,6 +24,19 @@ class Net(torch.nn.Module):
         return x
 
 
+class LSTMNet(torch.nn.Module):
+    def __init__(self, input_size, hidden_size=256, output_size=1, num_layers=1):
+        super(LSTMNet, self).__init__()
+
+        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers)
+        self.fc = torch.nn.Linear(hidden_size, output_size)
+
+    def forward(self, x):
+        x, _ = self.lstm(x)  # x is input with size (seq_len, batch_size, input_size)
+        x = self.fc(x)
+        return x[-1, :, ]
+
+
 def selective_initialization(component_points, population_size, machine_number):
     # assignment_result = [[0 for _ in range(len(component_points))] for _ in range(machine_number)]
     assignment_result = []
@@ -44,58 +59,74 @@ def optimizer_hyperheuristc(pcb_data, component_data, machine_number):
 if __name__ == '__main__':
     warnings.simplefilter(action='ignore', category=FutureWarning)
 
-    train_file, test_file = 'train_data.txt', 'test_data.txt'
+    parser = argparse.ArgumentParser(description='network training implementation')
-    num_epochs = 30000
+    parser.add_argument('--train', default=True, type=bool, help='determine whether training the network')
+    parser.add_argument('--save', default=True, type=bool,
+                        help='determine whether saving the parameters of network, linear regression model, etc.')
+    parser.add_argument('--overwrite', default=False, type=bool,
+                        help='determine whether overwriting the training and testing data')
+    parser.add_argument('--train_file', default='train_data.txt', type=str, help='training file path')
+    parser.add_argument('--test_file', default='test_data.txt', type=str, help='testing file path')
+    parser.add_argument('--num_epochs', default=15000, type=int, help='number of epochs for training process')
+    parser.add_argument('--batch_size', default=100000, type=int, help='size of training batch')
+    parser.add_argument('--lr', default=1e-4, type=float, help='learning rate for the network')
+
+    params = parser.parse_args()
 
     data_mgr = DataMgr()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # batch_size = 40000
+    if params.overwrite:
-    # for _ in range(batch_size):
+        file = {params.train_file: params.batch_size,
-    #     pcb_data, component_data = data_mgr.generator()     # random generate a PCB data
+                params.test_file: params.batch_size // data_mgr.get_update_round() // 5}
-    #     # data_mgr.remover()                                # 移除最近一次保存数据
+        for file_name, file_batch_size in file.items():
-    #     # data_mgr.saver('data/' + train_file, pcb_data)     # 保存新数据
+            for _ in range(int(file_batch_size)):
-    #
+                with open('opt/' + file_name, 'a') as f:
-    #     info = base_optimizer(1, pcb_data, component_data,
+                    mode = file_name.split('.')[0].split('_')[0]
-    #                           feeder_data=pd.DataFrame(columns=['slot', 'part', 'arg']), method='feeder_scan',
+                    pcb_data, component_data = data_mgr.generator(mode)     # random generate a PCB data
-    #                           hinter=True)
+                    # data_mgr.remover()                                # remove the last saved data
-    #
+                    # data_mgr.saver('data/' + file_name, pcb_data)     # save new data
-    #     data_mgr.recorder('opt/' + train_file, info, pcb_data, component_data)
 
-    train, save = True, True
+                    info = base_optimizer(1, pcb_data, component_data,
-    learning_rate = 0.0005
+                                          feeder_data=pd.DataFrame(columns=['slot', 'part', 'arg']),
+                                          method='feeder_scan',
+                                          hinter=True)
+
+                    data_mgr.recorder(f, info, pcb_data, component_data)
+                f.close()
 
     net = Net(input_size=data_mgr.get_feature(), output_size=1).to(device)
-    if train:
+    if params.train:
-        data = data_mgr.loader('opt/' + train_file)
+        data = data_mgr.loader('opt/' + params.train_file)
         x_fit, y_fit = np.array(data[2:]).T, np.array([data[1]]).T
         lr = LinearRegression()
         lr.fit(x_fit, y_fit)
 
-        x_train, y_train = np.array(data[0]), lr.predict(x_fit)     # np.array(data[1])
+        x_train, y_train = np.array(data[0][::10]), lr.predict(x_fit[::10])
+        # x_train, y_train = np.array(data[0]), np.array(data[2])
+
         x_train = torch.from_numpy(x_train.reshape((-1, np.shape(x_train)[1]))).float().to(device)
         y_train = torch.from_numpy(y_train.reshape((-1, 1))).float().to(device)
 
-        optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
+        optimizer = torch.optim.Adam(net.parameters(), lr=params.lr)
-        # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1000, gamma=0.1)
+        # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=6000, gamma=0.8)
 
         loss_func = torch.nn.MSELoss()
 
-        for epoch in range(num_epochs):
+        for epoch in range(params.num_epochs):
             pred = net(x_train)
             loss = loss_func(pred, y_train)
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
             # scheduler.step()
-            if epoch % 200 == 0:
+            if epoch % 50 == 0:
                 print('Epoch:  ', epoch, ', Loss: ', loss.item())
                 if loss.item() < 1e-4:
                     break
 
         net_predict = net(x_train).view(-1)
-        # pred_time, real_time = net_predict.cpu().detach().numpy(), y_train.view(-1).cpu().detach().numpy()
+        pred_time, real_time = net_predict.cpu().detach().numpy(), y_train.view(-1).cpu().detach().numpy()
-        pred_time, real_time = net_predict.cpu().detach().numpy(), np.array(data[1])
 
         pred_error = np.array([])
         for t1, t2 in np.nditer([pred_time, real_time]):
@@ -107,21 +138,24 @@ if __name__ == '__main__':
 
         mse = np.linalg.norm((net_predict - y_train.view(-1)).cpu().detach().numpy())
         print(f'mean square error for training data result : {mse: 2f} ')
-        if save:
+        if params.save:
-            torch.save(net.state_dict(), 'model_state.pth')
+            if not os.path.exists('model'):
-            with open('lr_model.pkl', 'wb') as f:
+                os.mkdir('model')
+            torch.save(net.state_dict(), 'model/net_model.pth')
+            with open('model/lr_model.pkl', 'wb') as f:
                 pickle.dump(lr, f)
-            joblib.dump(lr, "lr_model.m")
+            # torch.save(optimizer.state_dict(), 'model/optimizer_state.pth')
-            # torch.save(optimizer.state_dict(), 'optimizer_state.pth')
     else:
-        with open('lr_model.pkl', 'rb') as f:
+        with open('model/lr_model.pkl', 'rb') as f:
             lr = pickle.load(f)
-        net.load_state_dict(torch.load('model_state.pth'))
+        net.load_state_dict(torch.load('model/net_model.pth'))
         # optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
-        # optimizer.load_state_dict(torch.load('optimizer_state.pth'))
+        # optimizer.load_state_dict(torch.load('model/optimizer_state.pth'))
 
-    data = data_mgr.loader('opt/' + test_file)
+    data = data_mgr.loader('opt/' + params.test_file)
+    # x_test, y_test = np.array(data[0]), np.array(data[1])
     x_test, y_test = np.array(data[0]), lr.predict(np.array(data[2:]).T)
+
     x_test, y_test = torch.from_numpy(x_test.reshape((-1, np.shape(x_test)[1]))).float().to(device), \
                      torch.from_numpy(y_test.reshape((-1, 1))).float().to(device)
 
@@ -133,7 +167,8 @@ if __name__ == '__main__':
         pred_error = np.array([])
         for t1, t2 in np.nditer([pred_time, real_time]):
             pred_error = np.append(pred_error, abs(t1 - t2) / (t2 + 1e-10) * 100)
-
+        print(pred_time)
+        print(real_time)
         print('--------------------------------------')
         print(f'average prediction error for test data : {np.average(pred_error): .2f}% ')
         print(f'maximum prediction error for test data : {np.max(pred_error): .2f}% ')