CODEWITHSUNDEEP
neural-networkpytorchautoencoder
doctopus
doctopus

Reputation: 5657

Extracting the encoded representations of an image in PyTorch?

I have a template autoencoder neural net model made in PyTorch that I'm using on the Omniglot dataset. I'd like to extract the encoded representations of the image, but I'm unsure how.

# 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)

Upvotes: 2

Views: 1100

Answers (1)

Harshit Kumar
Harshit Kumar

Reputation: 12837

You can simply return the encoded output in the forward function as follows:

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

Modify the training function a little bit:

output, encoded_output = model(img)

OR you can simply call encoder:

encoded_output = model.encoder(img)

Upvotes: 3

Related Questions