Simulation_Regression.py

import numpy as np


import torch
import torch.nn as nn
import torch.nn.functional as F

import pandas as pd
import argparse
import errno

import os
from mask_module import Mask_Linear, Mask_Conv2d, Mask_BatchNorm2d, Mask_ReLU, Mask_AvgPool2d, Mask_Sequential


from collections import OrderedDict, namedtuple


parser = argparse.ArgumentParser(description='Simulation Regression')

# Basic Setting
parser.add_argument('--data_index', default=1, type = int, help = 'set data index')
args = parser.parse_args()

class my_Net(torch.nn.Module):
    def __init__(self):
        super(my_Net, self).__init__()
        self.fc1 = Mask_Linear(2000, 6)
        self.fc2 = Mask_Linear(6,4)
        self.fc3 = Mask_Linear(4,3)
        self.fc4 = Mask_Linear(3,1)
        self._masks = OrderedDict()

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        return x

    def named_masks(self, prefix='', recurse=True):
        gen = self._named_members(
            lambda module: module._masks.items(),
            prefix=prefix, recurse=recurse)
        for elem in gen:
            yield elem

    def to(self, *args, **kwargs):
        super(my_Net, self).to(*args, **kwargs)
        device, dtype, non_blocking = torch._C._nn._parse_to(*args, **kwargs)
        for name, para in self.named_masks():
            para.data = para.to(device)


    def update_mask(self, user_mask):
        for name, para in self.named_masks():
            para.data = user_mask[name]

    def get_mask(self):
        total_mask = {}
        for name, para in self.named_masks():
            total_mask[name] = para.clone()
        return total_mask


def main():
    data_index = args.data_index
    subn = 500


    NTrain = 10000
    Nval = 1000
    NTest = 1000
    TotalP = 2000

    x_train = np.matrix(np.zeros([NTrain, TotalP]))
    y_train = np.matrix(np.zeros([NTrain, 1]))

    x_val = np.matrix(np.zeros([Nval, TotalP]))
    y_val = np.matrix(np.zeros([Nval, 1]))

    x_test = np.matrix(np.zeros([NTest, TotalP]))
    y_test = np.matrix(np.zeros([NTest, 1]))

    temp = np.matrix(pd.read_csv("./data/regression/" + str(data_index) + "/x_train.csv"))
    x_train[:, :] = temp[:, 1:]
    temp = np.matrix(pd.read_csv("./data/regression/" + str(data_index) + "/y_train.csv"))
    y_train[:, :] = temp[:, 1:]
    temp = np.matrix(pd.read_csv("./data/regression/" + str(data_index) + "/x_val.csv"))
    x_val[:, :] = temp[:, 1:]
    temp = np.matrix(pd.read_csv("./data/regression/" + str(data_index) + "/y_val.csv"))
    y_val[:, :] = temp[:, 1:]
    temp = np.matrix(pd.read_csv("./data/regression/" + str(data_index) + "/x_test.csv"))
    x_test[:, :] = temp[:, 1:]
    temp = np.matrix(pd.read_csv("./data/regression/" + str(data_index) + "/y_test.csv"))
    y_test[:, :] = temp[:, 1:]
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


    x_train = torch.FloatTensor(x_train).to(device)
    y_train = torch.FloatTensor(y_train).to(device)
    x_val = torch.FloatTensor(x_val).to(device)
    y_val = torch.FloatTensor(y_val).to(device)
    x_test = torch.FloatTensor(x_test).to(device)
    y_test = torch.FloatTensor(y_test).to(device)


    np.random.seed(data_index)
    torch.manual_seed(data_index)

    net = my_Net()
    net.to(device)
    loss_func = nn.MSELoss()

    step_lr = 0.005

    step_lr = step_lr/NTrain
    optimization = torch.optim.SGD(net.parameters(), lr=step_lr)

    sigma = torch.FloatTensor([1]).to(device)
    # sigma.requires_grad = True

    max_loop = 100001
    PATH = './result/regression/'

    if not os.path.isdir(PATH):
        try:
            os.makedirs(PATH)
        except OSError as exc:  # Python >2.5
            if exc.errno == errno.EEXIST and os.path.isdir(PATH):
                pass
            else:
                raise

    show_information = 100

    para_path = []
    para_gamma_path = []
    for para in net.parameters():
        para_path.append(np.zeros([max_loop // show_information + 1] + list(para.shape)))
        para_gamma_path.append(np.zeros([max_loop // show_information + 1] + list(para.shape)))

    train_loss_path = np.zeros([max_loop // show_information + 1])
    val_loss_path = np.zeros([max_loop // show_information + 1])
    test_loss_path = np.zeros([max_loop // show_information + 1])


    MH_loop = 5
    lambda_n = 0.01
    prior_sigma = 0.01

    prior_sigma_0 = 0.001

    proposal_a = 1000
    proposal_b = 100

    current_ones = 0
    current_log_proposal = 0
    current_mask = {}
    new_mask = {}
    for name, para in net.named_parameters():
        probability = para.abs().mul(-proposal_b).exp().mul(proposal_a).add(1).pow(-1)
        current_ones = current_ones + para.numel()
        current_log_proposal = current_log_proposal + probability.log().sum()
        new_mask[name] = torch.ones_like(para)
        current_mask[name] = torch.ones_like(para)

    temperature = 0.01


    for iter in range(max_loop):
        if subn == NTrain:
            subsample = range(NTrain)
        else:
            subsample = np.random.choice(range(NTrain), size=subn, replace=False)

        net.zero_grad()

        if iter < 5000:
            proposal_b = 1000
            prior_sigma_0 = 0.01
        if iter >=5000 and iter < 10000:
            proposal_b = 400
            prior_sigma_0 = 0.001
        if iter >=10000 and iter < 15000:
            proposal_b = 200
        if iter >= 15000 and iter < 20000:
            proposal_b = 100
        if iter >= 20000:
            proposal_b = 50



        for MH_iter in range(MH_loop):

            output = net(x_train[subsample,])
            loss = loss_func(output, y_train[subsample,])
            loss = loss.div(2 * sigma).add(sigma.log().mul(0.5))
            prior = 0
            for name, para in net.named_parameters():
                prior = prior + para.mul(current_mask[name]).pow(2).div(-2*prior_sigma).sum() + para.mul(1-current_mask[name]).pow(2).div(-2*prior_sigma_0).sum() \
                        + current_mask[name].sum().mul(-0.5*np.log(prior_sigma)) + (1-current_mask[name]).sum().mul(-0.5*np.log(prior_sigma_0))
            current_target = loss.mul(-NTrain).add(prior)
            current_loss = current_target.div(-MH_loop)


            new_ones = 0
            new_log_proposal = 0
            for name, para in net.named_parameters():
                probability = para.abs().mul(-proposal_b).exp().mul(proposal_a).add(1).pow(-1)
                new_mask[name] = torch.where(torch.rand_like(para) < probability, torch.ones_like(para), torch.zeros_like(para))
                new_ones = new_ones + new_mask[name].sum().item()
                new_log_proposal = new_log_proposal + new_mask[name].add(probability-1).abs().log().sum()
            net.update_mask(new_mask)

            new_output = net(x_train[subsample,])
            new_loss = loss_func(new_output, y_train[subsample,])
            new_loss = new_loss.div(2 * sigma).add(sigma.log().mul(0.5))
            new_prior = 0
            for name, para in net.named_parameters():
                new_prior = new_prior + para.mul(new_mask[name]).pow(2).div(-2*prior_sigma).sum() + para.mul(1-new_mask[name]).pow(2).div(-2*prior_sigma_0).sum() \
                        + new_mask[name].sum().mul(-0.5*np.log(prior_sigma)) + (1-new_mask[name]).sum().mul(-0.5*np.log(prior_sigma_0))
            new__target = new_loss.mul(-NTrain).add(new_prior)
            new_loss = new__target.div(-MH_loop)


            with torch.no_grad():
                log_MH_ratio = ((new__target - current_target)/temperature + (new_ones - current_ones)/temperature * np.log(lambda_n/(1-lambda_n)) + (current_log_proposal - new_log_proposal)).item()

            test_stat = np.random.uniform(0,1,1)
            if np.log(test_stat) < log_MH_ratio:
                new_loss.backward()
                with torch.no_grad():
                    current_log_proposal = new_log_proposal
                    current_ones = new_ones
                    temp = current_mask
                    current_mask = new_mask
                    new_mask = temp
            else:
                current_loss.backward()
                net.update_mask(current_mask)


        optimization.step()
        with torch.no_grad():
            for para in net.parameters():
                para.data = para + torch.FloatTensor(para.shape).to(device).normal_().mul(np.sqrt(step_lr * temperature))

        if iter % show_information == 0:
            print('iteration:', iter)
            with torch.no_grad():
                output = net(x_train)
                loss = loss_func(output, y_train)
                print("train loss:", loss)
                train_loss_path[iter // show_information] = loss.cpu().data.numpy()
                output = net(x_val)
                loss = loss_func(output, y_val)
                print("val loss:", loss)
                val_loss_path[iter // show_information] = loss.cpu().data.numpy()
                output = net(x_test)
                loss = loss_func(output, y_test)
                print("test loss:", loss)
                test_loss_path[iter // show_information] = loss.cpu().data.numpy()
                print('sigma:', sigma)

                for i, para in enumerate(net.parameters()):
                    para_path[i][iter // show_information,] = para.cpu().data.numpy()

                for i, (name,mask) in enumerate(net.named_masks()):
                    para_gamma_path[i][iter // show_information,] = mask.cpu().data.numpy()

                print('number of 1:', np.sum((np.max(para_gamma_path[0][0:(iter // show_information + 1),].mean(0) > 0.5, 0) > 0)))
                print('number of true:',
                      np.sum((np.max(para_gamma_path[0][0:(iter // show_information + 1),].mean(0) > 0.5, 0) > 0)[0:5]))

    import pickle

    filename = PATH + 'data_' + str(data_index) + "result.txt"
    f = open(filename, 'wb')
    pickle.dump([para_path, para_gamma_path, train_loss_path, val_loss_path, test_loss_path], f)
    f.close()


if __name__ == '__main__':
    main()