smt-optimizer/optimizer_hyperheuristic.py

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} ')