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]