Spaces:
Runtime error
Runtime error
File size: 4,570 Bytes
4d85df4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
"""
Most of the code in this file is taken from https://github.com/waterljwant/SSC/blob/master/models/DDR.py
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class SimpleRB(nn.Module):
def __init__(self, in_channel, norm_layer, bn_momentum):
super(SimpleRB, self).__init__()
self.path = nn.Sequential(
nn.Conv3d(in_channel, in_channel, kernel_size=3, padding=1, bias=False),
norm_layer(in_channel, momentum=bn_momentum),
nn.ReLU(),
nn.Conv3d(in_channel, in_channel, kernel_size=3, padding=1, bias=False),
norm_layer(in_channel, momentum=bn_momentum),
)
self.relu = nn.ReLU()
def forward(self, x):
residual = x
conv_path = self.path(x)
out = residual + conv_path
out = self.relu(out)
return out
"""
3D Residual Block，3x3x3 conv ==> 3 smaller 3D conv, refered from DDRNet
"""
class Bottleneck3D(nn.Module):
def __init__(
self,
inplanes,
planes,
norm_layer,
stride=1,
dilation=[1, 1, 1],
expansion=4,
downsample=None,
fist_dilation=1,
multi_grid=1,
bn_momentum=0.0003,
):
super(Bottleneck3D, self).__init__()
# often，planes = inplanes // 4
self.expansion = expansion
self.conv1 = nn.Conv3d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = norm_layer(planes, momentum=bn_momentum)
self.conv2 = nn.Conv3d(
planes,
planes,
kernel_size=(1, 1, 3),
stride=(1, 1, stride),
dilation=(1, 1, dilation[0]),
padding=(0, 0, dilation[0]),
bias=False,
)
self.bn2 = norm_layer(planes, momentum=bn_momentum)
self.conv3 = nn.Conv3d(
planes,
planes,
kernel_size=(1, 3, 1),
stride=(1, stride, 1),
dilation=(1, dilation[1], 1),
padding=(0, dilation[1], 0),
bias=False,
)
self.bn3 = norm_layer(planes, momentum=bn_momentum)
self.conv4 = nn.Conv3d(
planes,
planes,
kernel_size=(3, 1, 1),
stride=(stride, 1, 1),
dilation=(dilation[2], 1, 1),
padding=(dilation[2], 0, 0),
bias=False,
)
self.bn4 = norm_layer(planes, momentum=bn_momentum)
self.conv5 = nn.Conv3d(
planes, planes * self.expansion, kernel_size=(1, 1, 1), bias=False
)
self.bn5 = norm_layer(planes * self.expansion, momentum=bn_momentum)
self.relu = nn.ReLU(inplace=False)
self.relu_inplace = nn.ReLU(inplace=True)
self.downsample = downsample
self.dilation = dilation
self.stride = stride
self.downsample2 = nn.Sequential(
nn.AvgPool3d(kernel_size=(1, stride, 1), stride=(1, stride, 1)),
nn.Conv3d(planes, planes, kernel_size=1, stride=1, bias=False),
norm_layer(planes, momentum=bn_momentum),
)
self.downsample3 = nn.Sequential(
nn.AvgPool3d(kernel_size=(stride, 1, 1), stride=(stride, 1, 1)),
nn.Conv3d(planes, planes, kernel_size=1, stride=1, bias=False),
norm_layer(planes, momentum=bn_momentum),
)
self.downsample4 = nn.Sequential(
nn.AvgPool3d(kernel_size=(stride, 1, 1), stride=(stride, 1, 1)),
nn.Conv3d(planes, planes, kernel_size=1, stride=1, bias=False),
norm_layer(planes, momentum=bn_momentum),
)
def forward(self, x):
residual = x
out1 = self.relu(self.bn1(self.conv1(x)))
out2 = self.bn2(self.conv2(out1))
out2_relu = self.relu(out2)
out3 = self.bn3(self.conv3(out2_relu))
if self.stride != 1:
out2 = self.downsample2(out2)
out3 = out3 + out2
out3_relu = self.relu(out3)
out4 = self.bn4(self.conv4(out3_relu))
if self.stride != 1:
out2 = self.downsample3(out2)
out3 = self.downsample4(out3)
out4 = out4 + out2 + out3
out4_relu = self.relu(out4)
out5 = self.bn5(self.conv5(out4_relu))
if self.downsample is not None:
residual = self.downsample(x)
out = out5 + residual
out_relu = self.relu(out)
return out_relu
|