| | import torch |
| | import torch.nn as nn |
| | from transformers import PreTrainedModel, ResNetBackbone |
| | from .configuration_conditional_unet import ConditionalUNetConfig |
| |
|
| | class UpSampleBlock(nn.Module): |
| | def __init__(self, in_channels, skip_channels, out_channels, condition_size): |
| | super(UpSampleBlock, self).__init__() |
| | self.up = nn.Upsample(scale_factor=2, mode='nearest') |
| | self.conv = nn.Sequential( |
| | nn.Conv2d(in_channels + skip_channels + condition_size, out_channels, kernel_size=3, padding=1), |
| | nn.BatchNorm2d(out_channels), |
| | nn.ReLU(inplace=True), |
| | nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), |
| | nn.BatchNorm2d(out_channels), |
| | nn.ReLU(inplace=True) |
| | ) |
| |
|
| | def forward(self, x, skip, condition, upsample=True): |
| | if upsample: |
| | x = self.up(x) |
| | b, _, h, w = x.size() |
| | |
| | condition = condition.view(b, -1, 1, 1).expand(-1, -1, h, w) |
| | x = torch.cat([x, skip, condition], dim=1) |
| | x = self.conv(x) |
| | return x |
| |
|
| | class ConditionalUNet(PreTrainedModel): |
| | config_class = ConditionalUNetConfig |
| |
|
| | def __init__(self, config): |
| | super().__init__(config) |
| | |
| | self.config = config |
| |
|
| | self.encoder_rep = config.encoder_rep |
| | self.encoder = ResNetBackbone.from_pretrained( |
| | self.encoder_rep, |
| | return_dict=False, |
| | output_hidden_states=True |
| | ) |
| | self.encoder.eval() |
| | self.encoder.requires_grad_(False) |
| |
|
| | self.num_labels = self.encoder.config.num_labels |
| | self.num_channels = self.encoder.config.num_channels |
| |
|
| | self.config.num_labels = self.num_labels |
| | self.config.num_channels = self.num_channels |
| |
|
| | hidden_sizes = self.encoder.config.hidden_sizes |
| | embedding_size = self.encoder.config.embedding_size |
| |
|
| | self.up_blocks = nn.ModuleList() |
| | num_stages = len(hidden_sizes) |
| |
|
| | in_channels = hidden_sizes[-1] |
| | for i in range(num_stages - 1, -1, -1): |
| | skip_channels = hidden_sizes[i - 1] if i > 0 else embedding_size |
| | out_channels = skip_channels |
| | self.up_blocks.append( |
| | UpSampleBlock( |
| | in_channels=in_channels, |
| | skip_channels=skip_channels, |
| | out_channels=out_channels, |
| | condition_size=self.num_labels |
| | ) |
| | ) |
| | in_channels = out_channels |
| |
|
| | self.final_conv = nn.Sequential( |
| | nn.Conv2d(in_channels + self.num_labels, in_channels, kernel_size=3, padding=1), |
| | nn.BatchNorm2d(in_channels), |
| | nn.ReLU(inplace=True), |
| | nn.Conv2d(in_channels, self.num_channels, kernel_size=1) |
| | ) |
| |
|
| | def forward(self, x, condition): |
| | outputs = self.encoder(x)[-1] |
| | x_stages = outputs[::-1] |
| | x = x_stages[0] |
| |
|
| | for i, up_block in enumerate(self.up_blocks): |
| | skip = x_stages[i + 1] if i + 1 < len(x_stages) else None |
| | upsample = i < len(self.up_blocks) - 1 |
| | if skip is not None: |
| | x = up_block(x, skip, condition, upsample=upsample) |
| | else: |
| | x = up_block(x, torch.zeros_like(x), condition, upsample=upsample) |
| |
|
| | x_upsampled = nn.functional.interpolate(x, scale_factor=4, mode='bilinear', align_corners=False) |
| | b, _, h, w = x_upsampled.size() |
| | condition_expanded = condition.view(b, -1, 1, 1).expand(-1, -1, h, w) |
| | final_input = torch.cat([x_upsampled, condition_expanded], dim=1) |
| | output = self.final_conv(final_input) |
| |
|
| | return output |
| |
|
