| 95 | |
| 96 | |
| 97 | class StackedDDCNNV2(nn.Module): |
| 98 | |
| 99 | def __init__(self, |
| 100 | in_filters, |
| 101 | n_blocks, |
| 102 | filters, |
| 103 | shortcut=True, |
| 104 | use_octave_conv=False, # not supported |
| 105 | pool_type="avg", |
| 106 | stochastic_depth_drop_prob=0.0): |
| 107 | super(StackedDDCNNV2, self).__init__() |
| 108 | |
| 109 | if use_octave_conv: |
| 110 | raise NotImplemented( |
| 111 | "Octave convolution not implemented in Pytorch version of Transnet!") |
| 112 | |
| 113 | assert pool_type == "max" or pool_type == "avg" |
| 114 | if use_octave_conv and pool_type == "max": |
| 115 | print( |
| 116 | "WARN: Octave convolution was designed with average pooling, not max pooling.") |
| 117 | |
| 118 | self.shortcut = shortcut |
| 119 | self.DDCNN = nn.ModuleList([ |
| 120 | DilatedDCNNV2(in_filters if i == 1 else filters * 4, filters, octave_conv=use_octave_conv, |
| 121 | activation=functional.relu if i != n_blocks else None) for i in range(1, n_blocks + 1) |
| 122 | ]) |
| 123 | self.pool = nn.MaxPool3d(kernel_size=( |
| 124 | 1, 2, 2)) if pool_type == "max" else nn.AvgPool3d(kernel_size=(1, 2, 2)) |
| 125 | self.stochastic_depth_drop_prob = stochastic_depth_drop_prob |
| 126 | |
| 127 | def forward(self, inputs): |
| 128 | x = inputs |
| 129 | shortcut = None |
| 130 | |
| 131 | for block in self.DDCNN: |
| 132 | x = block(x) |
| 133 | if shortcut is None: |
| 134 | shortcut = x |
| 135 | |
| 136 | x = functional.relu(x) |
| 137 | |
| 138 | if self.shortcut is not None: |
| 139 | if self.stochastic_depth_drop_prob != 0.: |
| 140 | if self.training: |
| 141 | if random.random() < self.stochastic_depth_drop_prob: |
| 142 | x = shortcut |
| 143 | else: |
| 144 | x = x + shortcut |
| 145 | else: |
| 146 | x = (1 - self.stochastic_depth_drop_prob) * x + shortcut |
| 147 | else: |
| 148 | x += shortcut |
| 149 | |
| 150 | x = self.pool(x) |
| 151 | return x |
| 152 | |
| 153 | |
| 154 | class DilatedDCNNV2(nn.Module): |