Как использовать пользовательский torch.autograd.Function в nn. Последовательная модель - PullRequest
Новая
       11

Как использовать пользовательский torch.autograd.Function в nn. Последовательная модель

1 голос
/ 09 апреля 2020

Можно ли использовать пользовательский torch.autograd.Function в объекте nn.Sequential или я должен явно использовать объект nn.Module с функцией forward. В частности, я пытаюсь реализовать разреженный авто-кодер, и мне нужно добавить L1-расстояние кода (скрытое представление) к потере. Я определил пользовательский torch.autograd.Function L1Penalty ниже, а затем попытался использовать его внутри nn.Sequential объекта, как показано ниже. Однако при запуске я получил ошибку TypeError: __main__.L1Penalty is not a Module subclass Как я могу решить эту проблему? 

class L1Penalty(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, l1weight = 0.1):
        ctx.save_for_backward(input)
        ctx.l1weight = l1weight
        return input, None

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_variables
        grad_input = input.clone().sign().mul(ctx.l1weight)
        grad_input+=grad_output
        return grad_input

model = nn.Sequential(
    nn.Linear(10, 10),
    nn.ReLU(),
    nn.Linear(10, 6),
    nn.ReLU(),
    # sparsity
    L1Penalty(),
    nn.Linear(6, 10),
    nn.ReLU(),
    nn.Linear(10, 10),
    nn.ReLU()
).to(device)

1 Ответ

1 голос
/ 10 апреля 2020

API nn.Module работает нормально, но вы не должны возвращать None в вашем методе L1Penalty forward. 

import torch, torch.nn as nn

class L1Penalty(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, l1weight = 0.1):
        ctx.save_for_backward(input)
        ctx.l1weight = l1weight
        return input

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_variables
        grad_input = input.clone().sign().mul(ctx.l1weight)
        grad_input+=grad_output
        return grad_input


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10,10)
        self.fc2 = nn.Linear(10,6)
        self.fc3 = nn.Linear(6,10)
        self.fc4 = nn.Linear(10,10)
        self.relu = nn.ReLU(inplace=True)
        self.penalty = L1Penalty()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.penalty.apply(x)
        x = self.fc3(x)
        x = self.relu(x)
        x = self.fc4(x)
        x = self.relu(x)
        return x


model = Model()
a = torch.rand(50,10)
b = model(a)
print(b.shape)
...