Извлечение закодированных представлений изображения в PyTorch? - PullRequest
Новая
       23

Извлечение закодированных представлений изображения в PyTorch?

2 голосов
/ 19 февраля 2020

У меня есть шаблон нейрокодера net, созданный в PyTorch, который я использую в наборе данных Omniglot. Я хотел бы извлечь закодированные представления изображения, но я не уверен, как. 

# Load data
mean = 0.5
std = 0.5
batch_size = 128
img_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((mean,), (std,))
])
dataset = Omniglot('.', download=True, transform=img_transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Define Autoencoder
class Autoencoder(nn.Module):

    def __init__(self, n=64):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(105*105, 256, bias=True),
            nn.ReLU(True),
            nn.Linear(256, 64, bias=True),
            nn.ReLU(True),
            nn.Linear(64, n, bias=True),
            nn.ReLU(True)
        )
        self.decoder = nn.Sequential(
            nn.Linear(n, 64, bias=True),
            nn.ReLU(True),
            nn.Linear(64, 256, bias=True),
            nn.ReLU(True),
            nn.Linear(256, 105*105, bias=True),
            nn.Tanh()
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

def train(num_epochs, dataloader, model, criterion, optimizer):
    for epoch in range(num_epochs):
        for data in dataloader:
            img, label = data
            img = img.view(img.size(0), -1)
            img = Variable(img).cuda()
            output = model(img)
            loss = criterion(output, img)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    return model

# Train model
num_epochs = 25
learning_rate = 1e-3
model = Autoencoder().cuda()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
trained_model = train(num_epochs, dataloader, model, criterion, optimizer)

1 Ответ

1 голос
/ 19 февраля 2020

Вы можете просто вернуть закодированный вывод в функции forward следующим образом: 

class Autoencoder(nn.Module):
...
    def forward(self, x):
        x = self.encoder(x)
        encoded_x = x
        x = self.decoder(x)
        return x, encoded_x

Немного изменить функцию обучения: 

output, encoded_output = model(img)

ИЛИ вы можете просто позвонить encoder: 

encoded_output = model.encoder(img)
...