<reponame>vishalbelsare/halp<gh_stars>1-10
import torch
import numpy as np
from torch.nn import Parameter
from halp.layers.embedding import BitCenterEmbedding, bit_center_embedding
from halp.utils.utils import void_cast_func, single_to_half_det, single_to_half_stoc
from unittest import TestCase
from halp.utils.test_utils import HalpTest
from torch.autograd.gradcheck import get_numerical_jacobian, iter_tensors, make_jacobian
from halp.layers.bit_center_layer_test import TestBitCenterLayer
import logging
import sys
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
logger = logging.getLogger()


class TestBitCenterEmbedding(TestBitCenterLayer, TestCase):
    '''
    Test the functionality of bit centering linear layers
    '''
    def get_config(self, type="grad_check"):
        config = {}
        if type == "grad_check":
            config["n_train_sample"] = 35
            config["num_embeddings"] = 17
            config["embedding_dim"] = 24
            config["bias"] = True
            config["cast_func"] = void_cast_func
            config["do_double"] = True
            config["seed"] = 0
            config["batch_size"] = 35
            config["seq_length"] = 11
        elif type == "fw_bw_proc":
            config["n_train_sample"] = 98
            config["num_embeddings"] = 13
            config["embedding_dim"] = 31
            config["bias"] = False
            config["cast_func"] = single_to_half_det
            config["do_double"] = False
            config["seed"] = 0
            config["batch_size"] = 33
            config["seq_length"] = 11
        else:
            raise Exception("Config type not supported!")
        return config

    def prepare_layer(self,
                      n_train_sample,
                      num_embeddings,
                      embedding_dim,
                      bias=False, # bias is dummy to adapt to the test interface
                      cast_func=void_cast_func,
                      do_double=True,
                      seed=0,
                      batch_size=1,
                      seq_length=10):
        layer = BitCenterEmbedding(
            num_embeddings=num_embeddings,
            embedding_dim=embedding_dim,
            cast_func=cast_func,
            n_train_sample=n_train_sample)
        # Note do_double = setup layer for gradient check, otherwise, it is for checking
        # the tensor properties
        if do_double:
            layer.double()
            layer.weight.data.copy_(
                torch.randn(
                    num_embeddings, embedding_dim, dtype=torch.double,
                    requires_grad=False).cuda())
            layer.weight_lp.data.copy_(layer.weight.data)
            layer.weight_delta.data.copy_(
                torch.randn(
                    num_embeddings, embedding_dim, dtype=torch.double,
                    requires_grad=True).cuda())
        layer.cuda()
        return layer

    def get_input(self,
                  n_train_sample,
                  num_embeddings,
                  embedding_dim,
                  bias=False,
                  cast_func=void_cast_func,
                  do_double=True,
                  seed=0,
                  batch_size=1,
                  seq_length=10):
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        self.target_dtype = None

        input_delta = Parameter(
            torch.zeros(n_train_sample, seq_length, dtype=torch.long).cuda().random_(0, num_embeddings),
            requires_grad=False)
        input_fp = Parameter(
            torch.zeros(n_train_sample, seq_length, dtype=torch.long).cuda().data.copy_(input_delta.data),
            requires_grad=False)
        return [
            input_fp,
        ], [
            input_delta,
        ]

    def check_layer_param_and_cache(self, layer):
        t_list = [(layer.weight, torch.float32, True, True),
                  (layer.weight_delta, torch.half, True, True),
                  (layer.weight_lp, torch.half, True, False),
                  (layer.input_cache, torch.long, False, False),
                  (layer.grad_output_cache, torch.half, False, False)]
        self.CheckLayerTensorProperty(t_list)
        self.CheckLayerTensorGradProperty(t_list)

    def get_numerical_grad(self,
                           layer,
                           input_fp,
                           input_delta,
                           perturb_eps,
                           target=None):
        grad_list = []
        layer.set_mode(do_offset=True)
        param_dict = layer.state_dict()
        # update the offset variable
        for name, param in layer.named_parameters():
            if name.endswith("_delta"):
                p_offset = param_dict[name.split("_delta")[0]]
                p_offset.data.add_(param)
        
        input_final = [  # for embedding the input_fp and input_delta are the same index array
            Parameter(x, requires_grad=False)
            for x, y in zip(input_fp, input_delta)
        ]
        output_final = layer(*input_final)
        loss = 0.5 * torch.sum(output_final**2)
        loss.backward()
        # use the gradient from 0.5*sum(output**2), this case has output = gradient wrt output
        num_input_grad = None
        grad_list.append(num_input_grad)

        def get_loss(x):
            output = layer(*x)
            return torch.sum(0.5 * output * output)

        grad_list += self.get_numerical_param_grad(layer, input_final, get_loss,
                                                   perturb_eps)
        return output_final, grad_list


if __name__ == "__main__":
    print(torch.__version__)
    unittest.main()
import torch
import numpy as np
from torch.nn import Parameter
from halp.layers.embedding import BitCenterEmbedding, bit_center_embedding
from halp.utils.utils import void_cast_func, single_to_half_det, single_to_half_stoc
from unittest import TestCase
from halp.utils.test_utils import HalpTest
from torch.autograd.gradcheck import get_numerical_jacobian, iter_tensors, make_jacobian
from halp.layers.bit_center_layer_test import TestBitCenterLayer
import logging
import sys
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
logger = logging.getLogger()


class TestBitCenterEmbedding(TestBitCenterLayer, TestCase):
    '''
    Test the functionality of bit centering linear layers
    '''
    def get_config(self, type="grad_check"):
        config = {}
        if type == "grad_check":
            config["n_train_sample"] = 35
            config["num_embeddings"] = 17
            config["embedding_dim"] = 24
            config["bias"] = True
            config["cast_func"] = void_cast_func
            config["do_double"] = True
            config["seed"] = 0
            config["batch_size"] = 35
            config["seq_length"] = 11
        elif type == "fw_bw_proc":
            config["n_train_sample"] = 98
            config["num_embeddings"] = 13
            config["embedding_dim"] = 31
            config["bias"] = False
            config["cast_func"] = single_to_half_det
            config["do_double"] = False
            config["seed"] = 0
            config["batch_size"] = 33
            config["seq_length"] = 11
        else:
            raise Exception("Config type not supported!")
        return config

    def prepare_layer(self,
                      n_train_sample,
                      num_embeddings,
                      embedding_dim,
                      bias=False, # bias is dummy to adapt to the test interface
                      cast_func=void_cast_func,
                      do_double=True,
                      seed=0,
                      batch_size=1,
                      seq_length=10):
        layer = BitCenterEmbedding(
            num_embeddings=num_embeddings,
            embedding_dim=embedding_dim,
            cast_func=cast_func,
            n_train_sample=n_train_sample)
        # Note do_double = setup layer for gradient check, otherwise, it is for checking
        # the tensor properties
        if do_double:
            layer.double()
            layer.weight.data.copy_(
                torch.randn(
                    num_embeddings, embedding_dim, dtype=torch.double,
                    requires_grad=False).cuda())
            layer.weight_lp.data.copy_(layer.weight.data)
            layer.weight_delta.data.copy_(
                torch.randn(
                    num_embeddings, embedding_dim, dtype=torch.double,
                    requires_grad=True).cuda())
        layer.cuda()
        return layer

    def get_input(self,
                  n_train_sample,
                  num_embeddings,
                  embedding_dim,
                  bias=False,
                  cast_func=void_cast_func,
                  do_double=True,
                  seed=0,
                  batch_size=1,
                  seq_length=10):
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        self.target_dtype = None

        input_delta = Parameter(
            torch.zeros(n_train_sample, seq_length, dtype=torch.long).cuda().random_(0, num_embeddings),
            requires_grad=False)
        input_fp = Parameter(
            torch.zeros(n_train_sample, seq_length, dtype=torch.long).cuda().data.copy_(input_delta.data),
            requires_grad=False)
        return [
            input_fp,
        ], [
            input_delta,
        ]

    def check_layer_param_and_cache(self, layer):
        t_list = [(layer.weight, torch.float32, True, True),
                  (layer.weight_delta, torch.half, True, True),
                  (layer.weight_lp, torch.half, True, False),
                  (layer.input_cache, torch.long, False, False),
                  (layer.grad_output_cache, torch.half, False, False)]
        self.CheckLayerTensorProperty(t_list)
        self.CheckLayerTensorGradProperty(t_list)

    def get_numerical_grad(self,
                           layer,
                           input_fp,
                           input_delta,
                           perturb_eps,
                           target=None):
        grad_list = []
        layer.set_mode(do_offset=True)
        param_dict = layer.state_dict()
        # update the offset variable
        for name, param in layer.named_parameters():
            if name.endswith("_delta"):
                p_offset = param_dict[name.split("_delta")[0]]
                p_offset.data.add_(param)
        
        input_final = [  # for embedding the input_fp and input_delta are the same index array
            Parameter(x, requires_grad=False)
            for x, y in zip(input_fp, input_delta)
        ]
        output_final = layer(*input_final)
        loss = 0.5 * torch.sum(output_final**2)
        loss.backward()
        # use the gradient from 0.5*sum(output**2), this case has output = gradient wrt output
        num_input_grad = None
        grad_list.append(num_input_grad)

        def get_loss(x):
            output = layer(*x)
            return torch.sum(0.5 * output * output)

        grad_list += self.get_numerical_param_grad(layer, input_final, get_loss,
                                                   perturb_eps)
        return output_final, grad_list


if __name__ == "__main__":
    print(torch.__version__)
    unittest.main()
Test the functionality of bit centering linear layers # bias is dummy to adapt to the test interface # Note do_double = setup layer for gradient check, otherwise, it is for checking # the tensor properties # update the offset variable # for embedding the input_fp and input_delta are the same index array # use the gradient from 0.5*sum(output**2), this case has output = gradient wrt output