garyzgao's picture
Uploading the file for model definition.
165f2cc verified
import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
from typing import Optional
class DecoderBlock(nn.Module):
def __init__(self, in_channels: int, skip_channels: int, out_channels: int):
super(DecoderBlock, self).__init__()
self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
self.conv1 = nn.Conv2d(out_channels + skip_channels, out_channels, kernel_size=3, padding=1)
self.bn1 = nn.BatchNorm2d(out_channels)
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
self.bn2 = nn.BatchNorm2d(out_channels)
self.relu2 = nn.ReLU(inplace=True)
def forward(self, x: torch.Tensor, skip: Optional[torch.Tensor]) -> torch.Tensor:
x = self.up(x)
if skip is not None:
if x.size() != skip.size():
x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=True)
x = torch.cat([x, skip], dim=1)
x = self.conv1(x)
x = self.bn1(x)
x = self.relu1(x)
x = self.conv2(x)
x = self.bn2(x)
x = self.relu2(x)
return x
class UNetTimmWithClassification(nn.Module):
def __init__(self,
encoder_name: str = 'resnet50',
encoder_weights: Optional[str] = 'imagenet',
num_classes_seg: int = 1,
num_classes_cls: int = 9
):
super(UNetTimmWithClassification, self).__init__()
self.encoder = timm.create_model(
encoder_name,
pretrained=(encoder_weights == 'imagenet'),
features_only=True,
in_chans=3
)
encoder_channels = self.encoder.feature_info.channels()
decoder_channels = [2048, 1024, 512, 256]
# decoder_channels = [512, 256, 128, 64]
self.decoder4 = DecoderBlock(
in_channels=decoder_channels[0],
skip_channels=encoder_channels[3],
out_channels=decoder_channels[1]
)
self.decoder3 = DecoderBlock(
in_channels=decoder_channels[1],
skip_channels=encoder_channels[2],
out_channels=decoder_channels[2]
)
self.decoder2 = DecoderBlock(
in_channels=decoder_channels[2],
skip_channels=encoder_channels[1],
out_channels=decoder_channels[3]
)
self.decoder1 = DecoderBlock(
in_channels=decoder_channels[3],
skip_channels=encoder_channels[0],
out_channels=decoder_channels[3]
)
self.final_up = nn.ConvTranspose2d(
in_channels=decoder_channels[-1],
out_channels=32,
kernel_size=2,
stride=2
)
self.final_conv_seg = nn.Conv2d(
in_channels=32,
out_channels=num_classes_seg,
kernel_size=1
)
#Cls head
self.classification_head = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(2048, 512),
# nn.Linear(encoder_channels[-1], 512),
nn.Dropout(0.2),
nn.Linear(512, 512),
nn.ReLU(),
nn.Linear(512, num_classes_cls)
)
if num_classes_cls > 1:
self.classification_activation = nn.Softmax(dim=1)
elif num_classes_cls == 1:
self.classification_activation = nn.Sigmoid()
else:
self.classification_activation = None
if self.classification_activation is not None:
self.classification_head.add_module("activation", self.classification_activation)
#Xavier weight initialize
if encoder_weights == 'xavier':
self.apply(self.xavier_init_weights)
def xavier_init_weights(self, m):
if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
nn.init.xavier_uniform_(m.weight)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.BatchNorm2d):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
def forward(self, x: torch.Tensor) -> tuple:
features = self.encoder(x)
C0, C1, C2, C3, C4 = features
cls = self.classification_head(C4)
D4 = self.decoder4(C4, C3)
D3 = self.decoder3(D4, C2)
D2 = self.decoder2(D3, C1)
D1 = self.decoder1(D2, C0)
x = self.final_up(D1)
seg = self.final_conv_seg(x)
return seg, cls