import os import pickle import numpy as np import torch.nn from base_optimizer.optimizer_interface import * from generator import * os.environ['KMP_DUPLICATE_LIB_OK'] = 'True' class Net(torch.nn.Module): def __init__(self, input_size, hidden_size=1024, output_size=1): super(Net, self).__init__() self.fc1 = torch.nn.Linear(input_size, hidden_size) self.relu = torch.nn.ReLU() # 激活函数 self.fc2 = torch.nn.Linear(hidden_size, output_size) def forward(self, x): x = self.fc1(x) x = self.relu(x) x = self.fc2(x) 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 = [] return assignment_result def optimizer_hyperheuristc(pcb_data, component_data, machine_number): # genetic-based hyper-heuristic crossover_rate, mutation_rate = 0.8, 0.1 population_size, n_generations = 200, 500 # todo: how to generate initial population (random?) # assignment_result = selective_initialization(component_points, population_size, machine_number) assignment_result = [] return assignment_result if __name__ == '__main__': warnings.simplefilter(action='ignore', category=FutureWarning) parser = argparse.ArgumentParser(description='network training implementation') 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") if params.overwrite: file = {params.train_file: params.batch_size, params.test_file: params.batch_size // data_mgr.get_update_round() // 5} for file_name, file_batch_size in file.items(): for _ in range(int(file_batch_size)): with open('opt/' + file_name, 'a') as f: mode = file_name.split('.')[0].split('_')[0] pcb_data, component_data = data_mgr.generator(mode) # random generate a PCB data # data_mgr.remover() # remove the last saved data # data_mgr.saver('data/' + file_name, pcb_data) # save new data info = base_optimizer(1, pcb_data, component_data, 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 params.train: 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][::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=params.lr) # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=6000, gamma=0.8) loss_func = torch.nn.MSELoss() 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 % 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_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('--------------------------------------') print(f'average prediction error for train data : {np.average(pred_error): .2f}% ') print(f'maximum prediction error for train data : {np.max(pred_error): .2f}% ') 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 params.save: if not os.path.exists('model'): 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) # torch.save(optimizer.state_dict(), 'model/optimizer_state.pth') else: with open('model/lr_model.pkl', 'rb') as f: lr = pickle.load(f) 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('model/optimizer_state.pth')) 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) net.eval() with torch.no_grad(): net_predict = net(x_test).view(-1) pred_time, real_time = net_predict.cpu().detach().numpy(), y_test.view(-1).cpu().detach().numpy() 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}% ') mse = np.linalg.norm(pred_time - real_time) print(f'mean square error for test data result : {mse: 2f} ')