给出输入是图像张量与超像素标签张量的超像素池化代码，并返回池化后的张量，pytorch实现

以下是一个简单的超像素池化代码示例，将输入的图像张量和超像素标签张量进行池化。假设输入的图像张量大小为 (batch_size, channels, height, width) ，超像素标签张量大小为 (batch_size, num_superpixels, height, width)，超像素池化大小为 pool_size。

import torch
import torch.nn.functional as F

class SuperpixelPooling(torch.nn.Module):
    def __init__(self, pool_size):
        super(SuperpixelPooling, self).__init__()
        self.pool_size = pool_size
        
    def forward(self, x, labels):
        # 输入张量维度为 (batch_size, channels, height, width)
        # 标签张量维度为 (batch_size, num_superpixels, height, width)
        
        batch_size, num_channels, height, width = x.size()
        num_superpixels = labels.size(1)
        
        # 将标签张量 (batch_size, num_superpixels, height, width) 变形为 (batch_size, num_superpixels, height*width)
        labels = labels.view(batch_size, num_superpixels, -1)
        
        # 计算每个超像素的大小
        sizes = torch.sum(labels, dim=2, keepdim=True)
        
        # 将输入张量 (batch_size, channels, height, width) 变形为 (batch_size, channels, height*width)
        x = x.view(batch_size, num_channels, -1)
        
        # 对输入张量和大小进行加权平均池化
        x = x * labels
        x = F.avg_pool1d(x, kernel_size=self.pool_size, stride=self.pool_size)
        sizes = F.avg_pool1d(sizes, kernel_size=self.pool_size, stride=self.pool_size)
        x = x / (sizes + 1e-8)
        
        # 将池化后的张量变形为 (batch_size, channels, height/pool_size, width/pool_size)
        x = x.view(batch_size, num_channels, height//self.pool_size, width//self.pool_size)
        
        return x

使用方法：

# 假设输入的图像张量为 x, 超像素标签张量为 labels，超像素池化大小为 pool_size
pooling_layer = SuperpixelPooling(pool_size)
pooled_x = pooling_layer(x, labels)