Trying to Implement SLIDE (Sub-LInear Deep learning Engine) CPU

username: its_sandy
cooked: <p>Hello all <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=9" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:"><br>
I’m trying to implement the algorithm described in <a href="https://arxiv.org/abs/1903.03129" rel="nofollow noopener">SLIDE : In Defense of Smart Algorithms over Hardware Acceleration for Large-Scale Deep Learning Systems</a> (<a href="https://github.com/keroro824/HashingDeepLearning" rel="nofollow noopener">https://github.com/keroro824/HashingDeepLearning</a>) as a standalone <em>torch.nn.module</em> layer. For a one-line description of the algo, it is a modified version of Sampled Softmax, where the sampled nodes for each input in the batch are different and are identified using hash functions. Weights undergo hogwild style updates. The algo is designed to use multithreading to decrease training time on the CPU.</p>
<p>I’ve tried 2 implementations for the sparse multiplication so far.</p>
<pre><code class="lang-auto">import concurrent.futures

class SlideLayer(nn.Module):
    def __init__(self, in_dim, out_dim):
        super(SlideLayer, self).__init__()
        self.linear = nn.Linear(in_dim, out_dim)

    def forward_1(self, in_values, active_out_indices):
        # in_values (batch x in_dim)   // dense input
        # active_out_indices (batch x num_active_out_nodes)
        # output_values (batch x num_active_out_nodes)  // sparse output
        # assume fixed num_active_out_nodes for each input in batch

        active_weights = self.linear.weight[[active_out_indices]]
        active_biases = self.linear.bias[[active_out_indices]]
        output_values = torch.bmm(active_weights, in_values.unsqueeze(-1)).squeeze(-1) + active_biases
        return output_values

    def forward_2(self, in_values, active_out_indices):
        batch_size = active_out_indices.shape[0]
        num_active_out_indices = active_out_indices.shape[1]

        def thread_fun(sample_num):
            torch.set_num_threads = 1
            cur_output_values= torch.stack([torch.dot(in_values[sample_num], self.linear.weight[i]) + self.linear.bias[i] for i in active_out_indices[sample_num]])
            output_values[sample_num] = cur_output_values

        output_values = [0]*batch_size
        with concurrent.futures.ThreadPoolExecutor(max_workers=min(64,batch_size)) as executor:
            executor.map(thread_fun, range(batch_size))
        return output_values
</code></pre>
<p>Profiling indicates that the main bottleneck for <strong>forward_1</strong> is the data copy because of advanced indexing. It requires 4x more time compared to the batch matrix multiplication (further the bmm by itself requires a small number of active nodes to prove meaningful).</p>
<p>I tried <strong>forward_2</strong> to avoid the data copy, but it is 1000x slower than forward_1. I assume this is because of GIL issues?</p>
<p>What is the best way to go about implementing this algorithm? Can this be done effectively in Python. Would <strong>distributed dataparallel</strong> help? Should I look into <strong>C++ extensions</strong> ?</p>
<p>Thank You<br>
<strong>Disclaimer: I’m a beginner in PyTorch and this is my first topic here. I apologise for the length of the post and any obvious mistakes.</strong></p>

username: ptrblck
cooked: <p>Your concern about the GIL might be valid (I haven’t verified it with your code though).<br>
Since your method seems to be similar to HOGWILD, you could have a look at the <a href="https://github.com/pytorch/examples/blob/master/mnist_hogwild/main.py">HOGWILD MNIST example</a>, which uses multiprocessing and would thus avoid the GIL. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>