binary_gru.py

class GRU(nn.Module):
    INPUT_SIZE = 603
    HIDDEN_SIZE = 100
    OUTPUT_SIZE = 1

    def __init__(self):
        super(GRU, self).__init__()
        self.gru = nn.GRU(self.INPUT_SIZE, self.HIDDEN_SIZE)
        self.linear = nn.Linear(self.HIDDEN_SIZE, self.OUTPUT_SIZE)
        self.sm = nn.Sigmoid()
        if torch.cuda.is_available():
            self.device = 'cuda'
        else:
            self.device = 'cpu'

    def forward(self, input, hidden):
        _, hn = self.gru(input, hidden)
        # reduce from 3 to 2 dimensions
        rearranged = hn.view(hn.size()[1], hn.size(2))
        out1 = self.linear(rearranged)
        out2 = self.sm(out1)
        return out2

    def initHidden(self, N):
        return Variable(torch.randn(1, N, self.HIDDEN_SIZE))


class CandidateDataset(Dataset):
    def __init__(self, x, y):
        self.len = x.shape[0]
        if torch.cuda.is_available():
            device = 'cuda'
        else:
            device = 'cpu'
        self.x_data = torch.as_tensor(x, device=device, dtype=torch.float)
        self.y_data = torch.as_tensor(y, device=device, dtype=torch.float)
    
    def __getitem__(self, index):
        return self.x_data[index], self.y_data[index]

    def __len__(self):
        return self.len


class NeuralModel():

    N_EPOCHS = 20
    BATCH_SIZE = 50
    INPUT_SIZE = 603
    
    def __init__(self, model):
        self.model = model
        self.optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
        self.loss_f = nn.BCELoss()
        if torch.cuda.is_available():
            self.device = 'cuda'
        else:
            self.device = 'cpu'

    def fit(self, candidate_count):
        candidate_ds = self.get_ds(candidate_count)
        train_loader = DataLoader(dataset = candidate_ds, batch_size = self.BATCH_SIZE, shuffle = True)
        self.model.train()
        for epoch in range(self.N_EPOCHS):
            running_loss = 0.0
            for batch_idx, (inputs, labels) in enumerate(train_loader):
                inputs, labels = Variable(inputs), Variable(labels)
                self.optimizer.zero_grad()
                inputs = inputs.view(-1, inputs.size()[0], self.INPUT_SIZE)
                # init hidden with number of rows in input
                y_pred = self.model(inputs, self.model.initHidden(inputs.size()[1]))
                loss = self.loss_f(y_pred, labels)
                loss.backward()
                self.optimizer.step()

                running_loss += loss.item()
                if batch_idx % 500 == 499:
                    print('[%d, %5d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 500))
                    running_loss = 0.0

def predict(self):
        feature_matrix = np.random.rand(10, self.INPUT_SIZE)
        x_data = torch.as_tensor(feature_matrix, device='cpu', dtype=torch.float)
        x_data = x_data.view(1, x_data.size()[0], self.INPUT_SIZE)
        y_pred = self.model(x_data, self.model.initHidden(x_data.size()[1]))
        # this gets max class and its energy for each candidate
        max_cand_class = torch.max(y_pred, 1)
        max_energy = 0
        max_index = -1
        for idx, class_label in enumerate(max_cand_class[1]):
            class_label = class_label.item()
            max_e = max_cand_class[0][idx].item()
            # choose test sample with max energy
            if class_label == 1 and max_e > max_energy:
                max_index = idx
                max_energy = max_e
        # this gets the index for the single candidate with max energy 
        values, indices = torch.max(max_cand_class[0], 0)
        # this gets the class (0 or 1) for the candidate with max energy
        max_energy_class = max_cand_class[1][max_index]

    def get_ds(self, candidate_count):
        feature_matrix = np.random.rand(candidate_count, self.INPUT_SIZE)
        target_matrix = np.zeros((candidate_count, 1), dtype=int)
        for i in range(candidate_count):
            if i % 5 == 0:
                target_matrix[i] = 1
        candidate_ds = CandidateDataset(feature_matrix, target_matrix)
        return candidate_ds
      
if __name__ == '__main__':
    gru = models.GRU()
    model = NeuralModel(gru)
    model.fit(5000)
    model.predict()