|
--- |
|
tags: |
|
- model_hub_mixin |
|
- pytorch_model_hub_mixin |
|
--- |
|
|
|
# Model Card: Time-Conditioned U-Net for MNIST |
|
|
|
## Model Details |
|
- **Architecture**: Time-Conditioned U-Net |
|
- **Dataset**: [Comic Faces Paired Synthetic](https://www.kaggle.com/datasets/defileroff/comic-faces-paired-synthetic) |
|
- **Batch Size**: 256 |
|
- **Image Size**: 28x28 |
|
- **Loss Function**: Mean Squared Error (MSE) |
|
- **Optimizer**: Adam (learning rate = 1e-4) |
|
|
|
## Model Architecture |
|
This model is a U-Net-based neural network that incorporates time conditioning using sinusoidal embeddings and an MLP. The architecture is designed for small grayscale images (e.g., MNIST) and consists of: |
|
|
|
### **Encoder (Contracting Path)**: |
|
- **Downsampling** using three `DoubleConv` layers with 32, 64, and 128 channels, respectively. |
|
- Time embedding added at each convolution block. |
|
- **Max pooling** used to reduce spatial dimensions. |
|
|
|
### **Decoder (Expanding Path)**: |
|
- **Upsampling** via bilinear interpolation. |
|
- Skip connections from encoder layers to corresponding decoder layers. |
|
- Two `DoubleConv` layers with 128+64 and 64+32 channels, respectively. |
|
- Final `1x1` convolution to map to the output. |
|
|
|
### **Time Embedding**: |
|
- Uses a sinusoidal positional encoding to represent timestep information. |
|
- An MLP refines the embedding before passing it to convolutional layers. |
|
|
|
## Implementation |
|
### **Generator (U-Net)** |
|
```python |
|
class UNet(nn.Module, PyTorchModelHubMixin): |
|
def __init__(self, in_channels=1, out_channels=1, time_embedding_dim=32): |
|
super(UNet, self).__init__() |
|
|
|
# Time embedding layer |
|
self.time_embedding = TimeEmbedding(time_embedding_dim) |
|
|
|
# Encoder |
|
self.down_conv1 = DoubleConv(in_channels, 32, time_embedding_dim) |
|
self.down_conv2 = DoubleConv(32, 64, time_embedding_dim) |
|
self.down_conv3 = DoubleConv(64, 128, time_embedding_dim) |
|
|
|
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2) |
|
self.upsample = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True) |
|
|
|
# Decoder |
|
self.up_conv2 = DoubleConv(128 + 64, 64, time_embedding_dim) |
|
self.up_conv1 = DoubleConv(64 + 32, 32, time_embedding_dim) |
|
self.final_conv = nn.Conv2d(32, out_channels, kernel_size=1) |
|
|
|
def forward(self, x, timesteps): |
|
t = self.time_embedding(timesteps) |
|
|
|
x1 = self.down_conv1(x, t) |
|
x2 = self.down_conv2(self.maxpool(x1), t) |
|
x3 = self.down_conv3(self.maxpool(x2), t) |
|
|
|
x = self.upsample(x3) |
|
x = torch.cat([x2, x], dim=1) |
|
x = self.up_conv2(x, t) |
|
|
|
x = self.upsample(x) |
|
x = torch.cat([x1, x], dim=1) |
|
x = self.up_conv1(x, t) |
|
|
|
return self.final_conv(x) |
|
``` |
|
Time Embedding |
|
```python |
|
class TimeEmbedding(nn.Module): |
|
def __init__(self, embedding_dim): |
|
super().__init__() |
|
self.mlp = nn.Sequential( |
|
nn.SiLU(), |
|
nn.Linear(embedding_dim, embedding_dim), |
|
) |
|
|
|
def forward(self, t): |
|
half_dim = self.embedding_dim // 2 |
|
embeddings = torch.exp(torch.arange(half_dim, device=t.device) * -(torch.log(torch.tensor(10000.0)) / (half_dim - 1))) |
|
embeddings = t[:, None] * embeddings[None, :] |
|
embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1) |
|
return self.mlp(embeddings) |
|
``` |
|
|
|
## Training Configuration |
|
- Batch Size: 256 |
|
- Image Size: 28x28 |
|
- Loss Function: Mean Squared Error (MSE) |
|
- Optimizer: Adam (learning rate = 1e-4) |
|
|
|
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] |
