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 141 142 143 144 145 146 147 148 149 150 151 152
| plain_conv_encoder_EAB.py import torch from torch import nn import numpy as np from typing import Union, Type, List, Tuple
from torch.nn.modules.conv import _ConvNd from torch.nn.modules.dropout import _DropoutNd from dynamic_network_architectures.building_blocks.simple_conv_blocks import StackedConvBlocks from dynamic_network_architectures.building_blocks.helper import maybe_convert_scalar_to_list, get_matching_pool_op, get_matching_convtransp from dynamic_network_architectures.building_blocks.all_attention import *
class PlainConvEncoder(nn.Module): has_shown_prompt = [False] # 类属性,用于记录是否已经显示过提示 def __init__(self, input_channels: int, n_stages: int, features_per_stage: Union[int, List[int], Tuple[int, ...]], conv_op: Type[_ConvNd], kernel_sizes: Union[int, List[int], Tuple[int, ...]], strides: Union[int, List[int], Tuple[int, ...]], n_conv_per_stage: Union[int, List[int], Tuple[int, ...]], conv_bias: bool = False, norm_op: Union[None, Type[nn.Module]] = None, norm_op_kwargs: dict = None, dropout_op: Union[None, Type[_DropoutNd]] = None, dropout_op_kwargs: dict = None, nonlin: Union[None, Type[torch.nn.Module]] = None, nonlin_kwargs: dict = None, return_skips: bool = False, nonlin_first: bool = False, pool: str = 'conv' ):
super().__init__() if isinstance(kernel_sizes, int): kernel_sizes = [kernel_sizes] * n_stages if isinstance(features_per_stage, int): features_per_stage = [features_per_stage] * n_stages if isinstance(n_conv_per_stage, int): n_conv_per_stage = [n_conv_per_stage] * n_stages if isinstance(strides, int): strides = [strides] * n_stages assert len(kernel_sizes) == n_stages, "kernel_sizes must have as many entries as we have resolution stages (n_stages)" assert len(n_conv_per_stage) == n_stages, "n_conv_per_stage must have as many entries as we have resolution stages (n_stages)" assert len(features_per_stage) == n_stages, "features_per_stage must have as many entries as we have resolution stages (n_stages)" assert len(strides) == n_stages, "strides must have as many entries as we have resolution stages (n_stages). " \ "Important: first entry is recommended to be 1, else we run strided conv drectly on the input"
stages = [] for s in range(n_stages): stage_modules = [] if pool == 'max' or pool == 'avg': if (isinstance(strides[s], int) and strides[s] != 1) or \ isinstance(strides[s], (tuple, list)) and any([i != 1 for i in strides[s]]): stage_modules.append(get_matching_pool_op(conv_op, pool_type=pool)(kernel_size=strides[s], stride=strides[s])) conv_stride = 1 elif pool == 'conv': conv_stride = strides[s] else: raise RuntimeError() stage_modules.append(StackedConvBlocks( n_conv_per_stage[s], conv_op, input_channels, features_per_stage[s], kernel_sizes[s], conv_stride, conv_bias, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, nonlin_first )) stages.append(nn.Sequential(*stage_modules)) input_channels = features_per_stage[s]
self.stages = nn.Sequential(*stages) self.output_channels = features_per_stage self.strides = [maybe_convert_scalar_to_list(conv_op, i) for i in strides] self.return_skips = return_skips
# we store some things that a potential decoder needs self.conv_op = conv_op self.norm_op = norm_op self.norm_op_kwargs = norm_op_kwargs self.nonlin = nonlin self.nonlin_kwargs = nonlin_kwargs self.dropout_op = dropout_op self.dropout_op_kwargs = dropout_op_kwargs self.conv_bias = conv_bias self.kernel_sizes = kernel_sizes
################################################## 边缘聚合块 ################################################## transpconv_op = get_matching_convtransp(conv_op=self.conv_op) self.downblock_channal = [32, 64, 128, 256, 512, 512] self.mattn = Spartial_Attention3d(kernel_size=3) self.mdcat1 = nn.Sequential( StackedConvBlocks( 1, conv_op, self.downblock_channal[0], self.downblock_channal[0], kernel_sizes[s], 2, conv_bias, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, nonlin_first)) self.mdcat2 = nn.Sequential( StackedConvBlocks( 1, conv_op, self.downblock_channal[0] + self.downblock_channal[1], self.downblock_channal[0] + self.downblock_channal[1], kernel_sizes[s], 2, conv_bias, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, nonlin_first)) self.mupcat3 = nn.Sequential( StackedConvBlocks( 1, conv_op, self.downblock_channal[0] + self.downblock_channal[1] + self.downblock_channal[2], self.downblock_channal[2], kernel_sizes[s], 1, conv_bias, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, nonlin_first)) self.gate3 = Gate(self.downblock_channal[2], self.downblock_channal[2]) self.mupcat2 = transpconv_op(self.downblock_channal[0] + self.downblock_channal[1] + self.downblock_channal[2], self.downblock_channal[1], kernel_size=2, stride=2, bias=False) self.gate2 = Gate(in_channels=self.downblock_channal[1], out_channels=self.downblock_channal[1]) self.mupcat1 = transpconv_op(self.downblock_channal[0] + self.downblock_channal[1] + self.downblock_channal[2], self.downblock_channal[0], kernel_size=4, stride=4, bias=False) self.gate1 = Gate(in_channels=self.downblock_channal[0], out_channels=self.downblock_channal[0]) ################################################## 边缘聚合块 ##################################################
def forward(self, x): ret = [] for s in self.stages: x = s(x) ret.append(x) if not PlainConvEncoder.has_shown_prompt[0]: # 如果还未显示过提示 print("################################################## EAB ##################################################") PlainConvEncoder.has_shown_prompt[0] = True # 将提示标记为已显示 # middle attention m1 = self.mattn(ret[0]) m2 = self.mattn(ret[1]) m3 = self.mattn(ret[2])
m1m2 = torch.cat([self.mdcat1(m1), m2], dim=1) # Shape : [B, C=32+64, D/2, H/2, W/2] m_feature = torch.cat([self.mdcat2(m1m2), m3], dim=1) # Shape : [B, C=32+64+128, D/4, H/4, W/4]
ret[0] = self.gate1(self.mupcat1(m_feature), ret[0]) ret[1] = self.gate2(self.mupcat2(m_feature), ret[1]) ret[2] = self.gate3(self.mupcat3(m_feature), ret[2])
''' tensors = {'m1': m1, 'm2': m2, 'm3': m3, 'm1m2': m1m2, 'm_feature': m_feature, 'gate_output1': gate_output1, 'gate_output2': gate_output2, 'self.mupcat3(m_feature)': self.mupcat3(m_feature)} for name, tensor in tensors.items(): print(f"Name: {name}, Shape: {tensor.shape}") '''
if self.return_skips: return ret else: return ret[-1]
def compute_conv_feature_map_size(self, input_size): output = np.int64(0) for s in range(len(self.stages)): if isinstance(self.stages[s], nn.Sequential): for sq in self.stages[s]: if hasattr(sq, 'compute_conv_feature_map_size'): output += self.stages[s][-1].compute_conv_feature_map_size(input_size) else: output += self.stages[s].compute_conv_feature_map_size(input_size) input_size = [i // j for i, j in zip(input_size, self.strides[s])] return output
|