---
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](http://yann.lecun.com/exdb/mnist/)  

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

```python
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](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
- Library: [More Information Needed]
- Docs: [More Information Needed]