metadata

tags:
  - model_hub_mixin
  - pytorch_model_hub_mixin
license: mit

ImageGenerationTAU: Autoencoder for MNIST Image Generation

Model Details

Model Architecture: Convolutional Autoencoder
Framework: PyTorch
Input Shape: (1, 28, 28) (Grayscale MNIST Images)
Latent Dimension: User-defined (hidden_dim)
Dataset: MNIST Handwritten Digits

Model Description

The ImageGenerationTAU model is a convolutional autoencoder designed for image generation and feature extraction from MNIST. It consists of:

An encoder that compresses the input image into a low-dimensional representation.
A decoder that reconstructs the original image from the compressed representation.

This model can be used for image denoising, feature learning, and generative tasks.

Training Details

Loss Function: Smooth L1 Loss
Optimizer: Adam
Batch Size: 512
Number of Epochs: TBD
Regularization: Batch Normalization

Model Architecture

class ImageGenerationTAU(nn.Module, PyTorchModelHubMixin):
    def __init__(self, hidden_dim):
        super(ImageGenerationTAU, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.Flatten(),
            nn.Linear(32 * 7 * 7, hidden_dim),
        )
        self.decoder = nn.Sequential(
            nn.Linear(hidden_dim, 32 * 7 * 7),
            nn.ReLU(),
            nn.Unflatten(1, (32, 7, 7)),
            nn.ConvTranspose2d(32, 64, kernel_size=2, stride=2),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.ConvTranspose2d(64, 1, kernel_size=2, stride=2),
            nn.Sigmoid(),
        )

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

This model has been pushed to the Hub using the PytorchModelHubMixin integration:

Library: [More Information Needed]
Docs: [More Information Needed]