Hyperdimensional-Computing/model.py at master · Cornell-RelaxML/Hyperdimensional-Computing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
#!/usr/bin/env python
# coding: utf-8
import torch
import numpy as np
import torch.nn.functional as F

'''
model for binary hdc
'''


class FastSign(torch.nn.Module):
    '''
    This is a fast version of the SignActivation.
    '''

    def __init__(self):
        super(FastSign, self).__init__()

    def forward(self, input):
        out_forward = torch.sign(input)
        out_backward = torch.clamp(input, -1.3, 1.3)
        return out_forward.detach() - out_backward.detach() + out_backward


class BinaryLinear(torch.nn.Linear):
    '''
    A fully connected layer with weights binarized to {-1, +1}.
    '''

    def __init__(self, in_features, out_features, bias=True):
        super(BinaryLinear, self).__init__(
            in_features, out_features, bias
        )
        self.binarize = FastSign()

    def forward(self, input):
        return F.linear(input, self.binarize(self.weight), self.bias)


class BModel(torch.nn.Module):
    def __init__(self, in_dim=32768, classes=10):
        super(BModel, self).__init__()
        self.in_dim = in_dim
        self.fc = BinaryLinear(self.in_dim, classes, bias=False)

    def forward(self, x):
        x = self.fc(x) * (1.0 / self.in_dim ** 0.5)
        return x


# '''
# model for cyclic group hdc of different order
# '''

class FastRound(torch.nn.Module):
    '''
    This is a fast version of the round.
    '''

    def __init__(self):
        super(FastRound, self).__init__()

    def forward(self, input):
        out_forward = torch.round(input)
        out_backward = input  # torch.clamp(input, -1.3, 1.3)
        return out_forward.detach() - out_backward.detach() + out_backward


class RoundLinear(torch.nn.Linear):
    '''
    A fully connected layer with weights rounded to closest integers
    '''

    def __init__(self, in_features, out_features, gorder, bias=True):
        super(RoundLinear, self).__init__(
            in_features, out_features, bias
        )
        self.gorder = gorder
        self.Bias = bias
        self.round = FastRound()
        self.radius = torch.nn.Parameter(torch.ones(1))  # 1.0

    def pts_map(self, x, r=1.0):
        theta = 2.0 * np.pi / (1.0 * self.gorder) * x
        pts = r * torch.stack([torch.cos(theta), torch.sin(theta)], -1)
        return pts

    def GroupSim(self, input, weight):
        map_weight = self.pts_map(weight, r=self.radius)
        map_input = self.pts_map(input).unsqueeze(1)
        return torch.sum(torch.sum(map_weight * map_input, dim=-1), dim=-1)

    def forward(self, input):
        weight_q = self.weight
        if not self.training:
            weight_q = self.round(self.weight)
        sims = self.GroupSim(input, weight_q)
        if self.Bias:
            sims += self.bias
        return sims


class GModel(torch.nn.Module):
    def __init__(self, gorder, in_dim=32768, classes=10):
        super(GModel, self).__init__()
        self.in_dim = in_dim
        self.fc = RoundLinear(self.in_dim, classes, gorder, bias=False)

    def forward(self, x):
        x = self.fc(x) * (1.0 / self.in_dim ** 0.5)
        return x