How can I speed up max pooling clusters of different sizes and shapes of an image?

Question

I have clustered the pixels of an image into clusters of different sizes and shapes. I want to max pool each cluster as fast as possible because the max pooling happens in one layer of my CNN.

To clarify: Input is a batch of images with the following shape [batch_size, height of image, width of image, number of channels]. I have clustered each image before I start training my CNN. So for each image I have a ndarray of labels with shape [height of image, width of image].

How can I max pool over all pixels of an image that have the same label for all labels? I understand how to do it with a of for loop but that is painstakingly slow. I am searching for a fast solution that ideally can max pool over every cluster of each image in less than a second.

For implementation, I use Python3.7 and PyTorch.

Answer 1

Yes, we can do it. Wrote sample implementation below. Hope it should help.

public class TestDummy {
    public static void main(String[] args) {
        Map<String, Integer> temp = new HashMap<>();
        temp.put("harish", 1);
        List<Map<String, Integer>> tempList = new ArrayList<>();
        tempList.add(temp);
        Map<String, List<Map<String, Integer>>> map = new HashMap<>();
        map.put("shyam", tempList);
        for (Map.Entry<String, List<Map<String, Integer>>> entry : map.entrySet()) {
            System.out.println("[Key] : " + entry.getKey() + " [Value] : " + entry.getValue());
        }
    }
}

Output will be as:

[Key] : shyam [Value] : [{harish=1}]

Answer 2

public static void main(String[] args) {
    Map<String, List<Map<String, Integer>>> map = new HashMap<>();
    map.put("Test", List.of(Map.of("TestValue", 1)));
    map.put("Test2", List.of(Map.of("TestValue2", 2)));

    Object[] entries = map.entrySet().toArray();
    for (int i = 0; i < entries.length; i++) {
        System.out.println(entries[i]);
    }
}