| --- | |
| tags: | |
| - unet | |
| - pix2pix | |
| library_name: pytorch | |
| --- | |
| # Pix2Pix UNet Model | |
| ## Model Description | |
| Custom UNet model for Pix2Pix image translation. | |
| - Image Size: 256 | |
| - Model Type: Small (256) | |
| ## Usage | |
| ```python | |
| import torch | |
| from small_256_model import UNet as small_UNet | |
| from big_1024_model import UNet as big_UNet | |
| # Load the model | |
| checkpoint = torch.load('model_weights.pth') | |
| model = big_UNet() if checkpoint['model_config']['big'] else small_UNet() | |
| model.load_state_dict(checkpoint['model_state_dict']) | |
| model.eval() | |
| Model Architecture | |
| UNet( | |
| (encoder): Sequential( | |
| (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) | |
| (1): ReLU(inplace=True) | |
| (2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) | |
| (3): ReLU(inplace=True) | |
| (4): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) | |
| (5): ReLU(inplace=True) | |
| ) | |
| (decoder): Sequential( | |
| (0): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) | |
| (1): ReLU(inplace=True) | |
| (2): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) | |
| (3): ReLU(inplace=True) | |
| (4): ConvTranspose2d(64, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) | |
| (5): Tanh() | |
| ) | |
| ) | |