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