Можно ли добавить обучаемый фильтр после автоэнкодера? - PullRequest
Новая
       4

Можно ли добавить обучаемый фильтр после автоэнкодера?

0 голосов
/ 02 марта 2020

Итак, я создаю шумоподавитель с помощью автоматического кодера. Идея состоит в том, что перед вычислением моей потери (после автоэнкодера) я применяю эмпирический фильтр Винера к текстурной карте изображения и добавляю его обратно в мой вывод автоэнкодера (добавляя обратно «потерянные детали»). Я закодировал этот фильтр с помощью PyTorch.

Моя первая попытка сработала, добавив фильтр в конец функции пересылки моего автоэнкодера. Я могу тренировать эту сеть, и она проходит через мой фильтр в процессе обучения. Однако, если я распечатываю свою сеть, фильтр не отображается в списке, и torchsummary не включает его при вычислении параметров.

Это заставляет меня думать, что я только обучаю автоэнкодеру, и мой фильтр фильтрует таким же образом каждый раз, а не учиться.

Возможно ли то, что я пытаюсь сделать?

Ниже мой автоэнкодер: 

class AutoEncoder(nn.Module):
"""Autoencoder simple implementation """
def __init__(self):
    super(AutoEncoder, self).__init__()
    # Encoder
    # conv layer
    self.block1 = nn.Sequential(
        nn.Conv2d(1, 48, 3, padding=1),
        nn.Conv2d(48, 48, 3, padding=1),
        nn.MaxPool2d(2),
        nn.BatchNorm2d(48),
        nn.LeakyReLU(0.1)

    )
    self.block2 = nn.Sequential(
        nn.Conv2d(48, 48, 3, padding=1),
        nn.MaxPool2d(2),
        nn.BatchNorm2d(48),
        nn.LeakyReLU(0.1)
    )
    self.block3 = nn.Sequential(
        nn.Conv2d(48, 48, 3, padding=1),
        nn.ConvTranspose2d(48, 48, 2, 2, output_padding=1),
        nn.BatchNorm2d(48),
        nn.LeakyReLU(0.1)
    )
    self.block4 = nn.Sequential(
        nn.Conv2d(96, 96, 3, padding=1),
        nn.Conv2d(96, 96, 3, padding=1),
        nn.ConvTranspose2d(96, 96, 2, 2),
        nn.BatchNorm2d(96),
        nn.LeakyReLU(0.1)
    )
    self.block5 = nn.Sequential(
        nn.Conv2d(144, 96, 3, padding=1),
        nn.Conv2d(96, 96, 3, padding=1),
        nn.ConvTranspose2d(96, 96, 2, 2),
        nn.BatchNorm2d(96),
        nn.LeakyReLU(0.1)
    )
    self.block6 = nn.Sequential(
        nn.Conv2d(97, 64, 3, padding=1),
        nn.BatchNorm2d(64),
        nn.Conv2d(64, 32, 3, padding=1),
        nn.BatchNorm2d(32),
        nn.Conv2d(32, 1, 3, padding=1),
        nn.LeakyReLU(0.1)
    )

    # self.blockNorm = nn.Sequential(
    #     nn.BatchNorm2d(1),
    #     nn.LeakyReLU(0.1)
    # )

def forward(self, x):
    # torch.autograd.set_detect_anomaly(True)
    # print("input: ", x.shape)
    pool1 = self.block1(x)
    # print("pool1: ", pool1.shape)
    pool2 = self.block2(pool1)
    # print("pool2: ", pool2.shape)
    pool3 = self.block2(pool2)
    # print("pool3: ", pool3.shape)
    pool4 = self.block2(pool3)
    # print("pool4: ", pool4.shape)
    pool5 = self.block2(pool4)
    # print("pool5: ", pool5.shape)
    upsample5 = self.block3(pool5)
    # print("upsample5: ", upsample5.shape)
    concat5 = torch.cat((upsample5, pool4), 1)
    # print("concat5: ", concat5.shape)
    upsample4 = self.block4(concat5)
    # print("upsample4: ", upsample4.shape)
    concat4 = torch.cat((upsample4, pool3), 1)
    # print("concat4: ", concat4.shape)
    upsample3 = self.block5(concat4)
    # print("upsample3: ", upsample3.shape)
    concat3 = torch.cat((upsample3, pool2), 1)
    # print("concat3: ", concat3.shape)
    upsample2 = self.block5(concat3)
    # print("upsample2: ", upsample2.shape)
    concat2 = torch.cat((upsample2, pool1), 1)
    # print("concat2: ", concat2.shape)
    upsample1 = self.block5(concat2)
    # print("upsample1: ", upsample1.shape)
    concat1 = torch.cat((upsample1, x), 1)
    # print("concat1: ", concat1.shape)
    output = self.block6(concat1)

    t_map = x - output

    for i in range(4):
        tensor = t_map[i, :, :, :]                 # Take each item in batch separately. Could account for this in Wiener instead

        tensor = torch.squeeze(tensor)              # Squeeze for Wiener input format

        tensor = wiener_3d(tensor, 0.05, 10)        # Apply Wiener with specified std and block size
        tensor = torch.unsqueeze(tensor, 0)         # unsqueeze to put back into block
        t_map[i, :, :, :] = tensor                  # put back into block

    filtered_output = output + t_map
    return filtered_output

для l oop на конец - применить фильтр к каждому изображению в пакете. Я понимаю, что это нельзя распараллелить, поэтому, если у кого-то есть идеи для этого, я был бы признателен за это. Я могу опубликовать функцию фильтра 'wiener 3d ()', если это поможет, просто хочу, чтобы публикация была короткой.

Я пытался определить собственный класс слоя с фильтром внутри него, но я очень заблудился быстро.

Любая помощь будет принята с благодарностью!

1 Ответ

0 голосов
/ 02 марта 2020

Если все, что вам нужно, это превратить ваш фильтр Wiener в модуль, выполните следующее: 

class WienerFilter(T.nn.Module):
    def __init__(self, param_a=0.05, param_b=10):
        super(WienerFilter, self).__init__()
        # This can be accessed like any other member via self.param_a
        self.register_parameter("param_a", T.nn.Parameter(T.tensor(param_a)))
        self.param_b = param_b

    def forward(self, input):
        for i in range(4):
            tensor = input[i]                
            tensor = torch.squeeze(tensor)
            tensor = wiener_3d(tensor, self.param_a, self.param_b)
            tensor = torch.unsqueeze(tensor, 0)
            input[i] = tensor 
        return input

Вы можете применить это, добавив строку 

self.wiener_filter = WienerFilter()

в Функция инициализации вашего AutoEncoder.

в форварде, затем вы все это замените для l oop на 

filtered_output = output + self.wiener_filter(t_map)

Torch знает, что модуль wiener_filter является модулем-членом, поэтому он будет перечислите модуль, если вы печатаете модули вашего AutoEncoder.

Если вы хотите распараллелить ваш фильтр Винера, вам нужно сделать это в терминах PyTorch, то есть использовать его операции над тензорами. Эти операции осуществляются параллельно.

...