import torch import torch.nn as nn import torch.nn.functional as F from .util import ResBlock3d, ResBlock3D_stage3, ResBlock2d, ResBlock3D_stage3_leak import torch import torch.nn as nn import torch.nn.functional as F # ------------------------------------------------------------------ # 1. 自适应共享权重的 2D 卷积 # ------------------------------------------------------------------ class RegionAwareAdaptiveNorm(nn.Module): def __init__(self, in_channels, eps=1e-8): """ 初始化 Region-Aware Adaptive Normalization 模块。 Args: - in_channels: 输入通道数。 - eps: 防止除零的小值。 """ super().__init__() self.in_channels = in_channels self.eps = eps # 可学习的 gamma 和 beta 分支 self.gamma_fc = nn.Sequential( nn.Conv2d(in_channels, in_channels, kernel_size=1, bias=True), nn.ReLU(inplace=True) ) self.beta_fc = nn.Sequential( nn.Conv2d(in_channels, in_channels, kernel_size=1, bias=True), nn.ReLU(inplace=True) ) def forward(self, F, mask): """ 应用 Region-Aware Adaptive Normalization。 Args: - F: 输入特征图，形状为 [N, C, H, W]。 - mask: 掩码，形状为 [N, 1, H, W]，值域在 [0, 1]。 Returns: - F_out: 区域自适应归一化后的特征图。 """ N, C, H, W = F.shape # 计算 mask 覆盖区域的统计量 F_masked = F * mask mean_masked = F_masked.sum(dim=(2, 3), keepdim=True) / (mask.sum(dim=(2, 3), keepdim=True) + self.eps) var_masked = ((F_masked - mean_masked) ** 2 * mask).sum(dim=(2, 3), keepdim=True) / (mask.sum(dim=(2, 3), keepdim=True) + self.eps) F_norm_masked = (F_masked - mean_masked) / torch.sqrt(var_masked + self.eps) # 计算 mask 外区域的统计量 complement_mask = 1 - mask F_complement = F * complement_mask mean_complement = F_complement.sum(dim=(2, 3), keepdim=True) / (complement_mask.sum(dim=(2, 3), keepdim=True) + self.eps) var_complement = ((F_complement - mean_complement) ** 2 * complement_mask).sum(dim=(2, 3), keepdim=True) / (complement_mask.sum(dim=(2, 3), keepdim=True) + self.eps) F_norm_complement = (F_complement - mean_complement) / torch.sqrt(var_complement + self.eps) # 通过 mask 外的区域学习 gamma 和 beta gamma = self.gamma_fc(F_complement) # [N, C, H, W] beta = self.beta_fc(F_complement) # [N, C, H, W] # 应用 gamma 和 beta 到归一化后的特征 F_norm_complement = F_norm_complement * gamma + beta # 融合两种区域的特征 F_out = mask * F_norm_masked + complement_mask * F_norm_complement return F_out class AdaptiveSharedWeightConv2d(nn.Module): def __init__( self, in_channels, out_channels, latent_size, kernel_size=3, stride=1, padding=1, bias=False, eps=1e-8, use_learned_mask=True, use_adaptive_norm=False ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,)*2 self.stride = stride self.padding = padding self.bias_flag = bias self.eps = eps # 单一卷积核权重 (共用) self.weight = nn.Parameter( torch.randn(out_channels, in_channels, *self.kernel_size) * 0.01 ) # 调制分支需要的 style_fc hidden_size = in_channels self.style_fc = nn.Sequential( nn.Linear(latent_size, hidden_size, bias=True), nn.LeakyReLU(negative_slope=0.2, inplace=True), nn.Linear(hidden_size, in_channels, bias=True) ) # bias，如果需要的话 if bias: self.bias_param = nn.Parameter(torch.zeros(out_channels)) else: self.bias_param = None # 可选：学习一个 mask (也可以外部喂 mask) self.use_learned_mask = use_learned_mask if use_learned_mask: self.mask_conv = nn.Sequential( nn.Conv2d(in_channels, 1, kernel_size=3, padding=1), nn.Sigmoid() ) else: self.mask_conv = None self.use_adaptive_norm = use_adaptive_norm if use_adaptive_norm: self.adaptive_norm = RegionAwareAdaptiveNorm(in_channels) def forward(self, x, latent, external_mask=None): """ x: [N, inC, H, W] latent: [N, latent_size] external_mask: [N, 1, H, W] (可选) returns: (out, mask) # mask 维度: [N,1,H,W] """ N, _, H, W = x.shape # 1) 标准卷积 out_std out_std = F.conv2d( x, self.weight, bias=None, stride=self.stride, padding=self.padding ) # 2) 调制卷积 out_mod style = self.style_fc(latent) # => [N, inC] style = style.unsqueeze(-1).unsqueeze(-1) # => [N, inC, 1, 1] w = self.weight.unsqueeze(0) # => [1, outC, inC, kH, kW] w_mod = w * style[:, None, :, :, :] # => [N, outC, inC, kH, kW] demod = torch.rsqrt((w_mod**2).sum(dim=(2,3,4), keepdim=True) + self.eps) w_mod = w_mod * demod # => [N, outC, inC, kH, kW] x_reshape = x.view(1, N*self.in_channels, H, W) w_mod_reshape = w_mod.view(N*self.out_channels, self.in_channels, *self.kernel_size) out_mod = F.conv2d( x_reshape, w_mod_reshape, bias=None, stride=self.stride, padding=self.padding, groups=N ) out_mod = out_mod.view(N, self.out_channels, out_mod.shape[2], out_mod.shape[3]) if self.bias_param is not None: out_mod = out_mod + self.bias_param.view(1, -1, 1, 1) # 3) 得到 mask if external_mask is not None: mask = external_mask elif self.mask_conv is not None: mask = self.mask_conv(x) # => [N,1,H,W] else: # 若无 mask_conv，也无 external_mask，就默认全用调制 => mask=1 mask = torch.ones_like(out_std[:,0:1,:,:]) # 4) 融合 # adaptive_intensity = mask * 0.1 # noise = torch.randn_like(mask, device = mask.device) * adaptive_intensity # mask = mask + noise out = mask * out_mod + (1 - mask) * out_std # out = out_mod # 5) 可选：应用自适应归一化 if self.use_adaptive_norm: out = self.adaptive_norm(out, mask) return out, mask # ------------------------------------------------------------------ # 2. 自适应共享权重的 3D 卷积 # ------------------------------------------------------------------ class AdaptiveSharedWeightConv3d(nn.Module): def __init__( self, in_channels, out_channels, latent_size, kernel_size=3, stride=1, padding=1, bias=False, eps=1e-8, use_learned_mask=True ): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,)*3 self.stride = stride if isinstance(stride, tuple) else (stride,)*3 self.padding = padding if isinstance(padding, tuple) else (padding,)*3 self.bias_flag = bias self.eps = eps # 单一3D卷积核 (共享) self.weight = nn.Parameter( torch.randn(out_channels, in_channels, *self.kernel_size) * 0.01 ) # style_fc: 将 latent => [N, inC] hidden_size = in_channels self.style_fc = nn.Sequential( nn.Linear(latent_size, hidden_size, bias=True), nn.LeakyReLU(negative_slope=0.2, inplace=True), nn.Linear(hidden_size, in_channels, bias=True) ) if bias: self.bias_param = nn.Parameter(torch.zeros(out_channels)) else: self.bias_param = None # 可选：学习一个空间 mask self.use_learned_mask = use_learned_mask if use_learned_mask: self.mask_conv = nn.Sequential( nn.Conv3d(in_channels, 1, kernel_size=3, padding=1), nn.Sigmoid() ) else: self.mask_conv = None def forward(self, x, latent, external_mask=None): """ x: [N, inC, D, H, W] latent: [N, latent_size] external_mask: [N, 1, D, H, W], 值在 [0,1] returns: (out, mask) # mask 维度: [N,1,D,H,W] """ N, _, D, H, W = x.shape # 1) 标准卷积 => out_std out_std = F.conv3d( x, self.weight, bias=None, stride=self.stride, padding=self.padding ) # 2) 调制卷积 => out_mod style = self.style_fc(latent) # => [N, inC] style = style.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) # => [N, inC, 1, 1, 1] w = self.weight.unsqueeze(0) # => [1, outC, inC, kD, kH, kW] w_mod = w * style[:, None, :, :, :, :] # => [N, outC, inC, kD, kH, kW] demod = torch.rsqrt((w_mod**2).sum(dim=(2,3,4,5), keepdim=True) + self.eps) w_mod = w_mod * demod x_reshape = x.view(1, N*self.in_channels, D, H, W) w_mod_reshape = w_mod.view(N*self.out_channels, self.in_channels, *self.kernel_size) out_mod = F.conv3d( x_reshape, w_mod_reshape, bias=None, stride=self.stride, padding=self.padding, groups=N ) D_out, H_out, W_out = out_mod.shape[2], out_mod.shape[3], out_mod.shape[4] out_mod = out_mod.view(N, self.out_channels, D_out, H_out, W_out) if self.bias_param is not None: out_mod = out_mod + self.bias_param.view(1, -1, 1, 1, 1) # 3) 得到 mask if external_mask is not None: mask = external_mask elif self.mask_conv is not None: mask = self.mask_conv(x) # => [N,1,D,H,W] else: mask = torch.ones_like(out_std[:,0:1,:,:,:]) # 4) 融合 => out out = mask * out_mod + (1 - mask) * out_std return out, mask # ------------------------------------------------------------------ # 3. 2D/3D 自适应残差块，分别返回 (out, mask_of_second_conv) # ------------------------------------------------------------------ class ResnetBlock_Adaptive2D(nn.Module): def __init__(self, dim=512, latent_size=512, activation=nn.ReLU(True), use_adaptive_norm=False): super().__init__() self.dim = dim self.act = activation self.conv1 = AdaptiveSharedWeightConv2d( in_channels=dim, out_channels=dim, latent_size=latent_size, kernel_size=3, padding=1, bias=True, use_learned_mask=True, use_adaptive_norm=use_adaptive_norm ) self.conv2 = AdaptiveSharedWeightConv2d( in_channels=dim, out_channels=dim, latent_size=latent_size, kernel_size=3, padding=1, bias=True, use_learned_mask=True, use_adaptive_norm=use_adaptive_norm ) def forward(self, x, dlatents_in_slice, dlatents2 = None): """ x: [N, dim, H, W] dlatents_in_slice: [N, latent_size] returns: (out, mask2) # mask2: [N,1,H,W] """ if dlatents2 is None: dlatents2 = dlatents_in_slice y, mask1 = self.conv1(x, dlatents2) y = self.act(y) y, mask2 = self.conv2(y, dlatents_in_slice) out = x + y return out, (mask1 + mask2)/2 class ResnetBlock_Adaptive3D(nn.Module): def __init__(self, dim=32, latent_size=512, activation=nn.ReLU(True)): super().__init__() self.dim = dim self.act = activation self.conv1 = AdaptiveSharedWeightConv3d( in_channels=dim, out_channels=dim, latent_size=latent_size, kernel_size=3, padding=1, bias=True, use_learned_mask=True ) self.conv2 = AdaptiveSharedWeightConv3d( in_channels=dim, out_channels=dim, latent_size=latent_size, kernel_size=3, padding=1, bias=True, use_learned_mask=True ) def forward(self, x, dlatents_in_slice): """ x: [N, dim, D, H, W] dlatents_in_slice: [N, latent_size] returns: (out, mask2) # mask2: [N,1,D,H,W] """ y, mask1 = self.conv1(x, dlatents_in_slice) y = self.act(y) y, mask2 = self.conv2(y, dlatents_in_slice) out = x + y return out, (mask1 + mask2)/2 # ------------------------------------------------------------------ # 5. 修改后的 transfer_model # - 在 forward 方法里添加 return_mask=False # - 如果 True，则收集每一层的 mask 并返回（3D mask 在 D 维做 mean 后得到 [N,1,H,W]） # - 如果 False，则只返回特征 # ------------------------------------------------------------------ class transfer_model(nn.Module): def __init__(self, latent_dim=512, n_blocks=4): super(transfer_model, self).__init__() activation = nn.ReLU(True) # 3D in BN_in = [] for i in range(3): BN_in += [ ResnetBlock_Adaptive3D(dim=32, latent_size=latent_dim, activation=activation) ] self.BottleNeck_3din = nn.Sequential(*BN_in) # 2D BN = [] for i in range(n_blocks): BN += [ ResnetBlock_Adaptive2D(dim=512, latent_size=latent_dim, activation=activation) ] self.BottleNeck_2d = nn.Sequential(*BN) # 3D out BN_out = [] for i in range(3): BN_out += [ ResnetBlock_Adaptive3D(dim=32, latent_size=latent_dim, activation=activation) ] self.BottleNeck_3dout = nn.Sequential(*BN_out) # 末端的普通 3D ResBlock (不产生 mask) self.resblocks_3d = nn.Sequential() for i in range(3): self.resblocks_3d.add_module( '3dr' + str(i), ResBlock3d(32, kernel_size=3, padding=1) ) def forward(self, x, dlatents, return_mask=False): """ x => [N, 32, D, H, W] dlatents => [N, latent_dim] return_mask => bool, True => 返回 (输出, [mask_2d_1, mask_2d_2, ...]) # 每层一个 [N,1,H,W] False => 只返回输出 (默认) """ # 如果需要收集 mask，则开一个 list mask_list = [] if return_mask else None # 1) 3D in for i in range(len(self.BottleNeck_3din)): x, mask_3d = self.BottleNeck_3din[i](x, dlatents) # => ( [N,32,D,H,W], [N,1,D,H,W] ) # 对 3D mask 在 D 维做 mean => [N,1,H,W] mask_2d = mask_3d.mean(dim=2, keepdim=False) # => [N,1,H,W] if return_mask: mask_list.append(mask_2d) # 2) reshape to 2D => [N, 32*D, H, W] bs, c, d, h, w = x.shape x = x.view(bs, c*d, h, w) # => [N, 32*D, H, W] # 2D blocks for i in range(len(self.BottleNeck_2d)): x, mask_2d = self.BottleNeck_2d[i](x, dlatents) # => ( [N,32*D,H,W], [N,1,H,W] ) if return_mask: mask_list.append(mask_2d) # reshape back => [N, 32, D, H, W] x = x.view(bs, c, d, h, w) # => [N, 32, D, H, W] # 3) 3D out for i in range(len(self.BottleNeck_3dout)): x, mask_3d = self.BottleNeck_3dout[i](x, dlatents) # => ( [N,32,D,H,W], [N,1,D,H,W] ) mask_2d = mask_3d.mean(dim=2, keepdim=False) # => [N,1,H,W] if return_mask: mask_list.append(mask_2d) # 4) 末端普通3D残差 (无mask) x = self.resblocks_3d(x) # => [N, 32, D, H, W] if return_mask: # 返回 (特征, [mask_2d_1, mask_2d_2, ...]) return x, mask_list else: # 只返回特征 return x class transfer_model2(nn.Module): def __init__(self, latent_dim=512, n_blocks=7): super(transfer_model2, self).__init__() activation = nn.ReLU(True) # # 3D in # BN_in = [] # for i in range(3): # BN_in += [ # ResnetBlock_Adaptive3D(dim=32, latent_size=latent_dim, activation=activation) # ] # self.BottleNeck_3din = nn.Sequential(*BN_in) # 2D BN = [] for i in range(n_blocks): BN += [ ResnetBlock_Adaptive2D(dim=512, latent_size=latent_dim, activation=activation) ] self.BottleNeck_2d = nn.Sequential(*BN) # 3D out # BN_out = [] # for i in range(3): # BN_out += [ # ResnetBlock_Adaptive3D(dim=32, latent_size=latent_dim, activation=activation) # ] # self.BottleNeck_3dout = nn.Sequential(*BN_out) # 末端的普通 3D ResBlock (不产生 mask) self.resblocks_3d = nn.Sequential() for i in range(6): self.resblocks_3d.add_module( '3dr' + str(i), ResBlock3d(32, kernel_size=3, padding=1) ) def forward(self, x, dlatents, return_mask=False): """ x => [N, 32, D, H, W] dlatents => [N, latent_dim] return_mask => bool, True => 返回 (输出, [mask_2d_1, mask_2d_2, ...]) # 每层一个 [N,1,H,W] False => 只返回输出 (默认) """ # 如果需要收集 mask，则开一个 list mask_list = [] if return_mask else None # 2) reshape to 2D => [N, 32*D, H, W] bs, c, d, h, w = x.shape x = x.view(bs, c*d, h, w) # => [N, 32*D, H, W] # 2D blocks for i in range(len(self.BottleNeck_2d)): x, mask_2d = self.BottleNeck_2d[i](x, dlatents) # => ( [N,32*D,H,W], [N,1,H,W] ) if return_mask: mask_list.append(mask_2d) # reshape back => [N, 32, D, H, W] x = x.view(bs, c, d, h, w) # => [N, 32, D, H, W] # 4) 末端普通3D残差 (无mask) x = self.resblocks_3d(x) # => [N, 32, D, H, W] if return_mask: # 返回 (特征, [mask_2d_1, mask_2d_2, ...]) return x, mask_list else: # 只返回特征 return x class transfer_model3(nn.Module): def __init__(self, latent_dim=512, n_blocks=7): super(transfer_model3, self).__init__() activation = nn.ReLU(True) # # 3D in # BN_in = [] # for i in range(3): # BN_in += [ # ResnetBlock_Adaptive3D(dim=32, latent_size=latent_dim, activation=activation) # ] # self.BottleNeck_3din = nn.Sequential(*BN_in) # 2D BN = [] for i in range(n_blocks): BN += [ ResnetBlock_Adaptive2D(dim=512, latent_size=latent_dim, activation=activation) ] self.BottleNeck_2d = nn.Sequential(*BN) # 3D out # BN_out = [] # for i in range(3): # BN_out += [ # ResnetBlock_Adaptive3D(dim=32, latent_size=latent_dim, activation=activation) # ] # self.BottleNeck_3dout = nn.Sequential(*BN_out) # 末端的普通 3D ResBlock (不产生 mask) self.resblocks_3d = nn.Sequential() for i in range(6): self.resblocks_3d.add_module( '3dr' + str(i), ResBlock3d(32, kernel_size=3, padding=1) ) self.fc = nn.Linear(256*7*7, 512) def forward(self, x, dlatents, dlatents2,return_mask=False): """ x => [N, 32, D, H, W] dlatents => [N, latent_dim] return_mask => bool, True => 返回 (输出, [mask_2d_1, mask_2d_2, ...]) # 每层一个 [N,1,H,W] False => 只返回输出 (默认) """ # 如果需要收集 mask，则开一个 list mask_list = [] if return_mask else None # 2) reshape to 2D => [N, 32*D, H, W] bs, c, d, h, w = x.shape x = x.view(bs, c*d, h, w) # => [N, 32*D, H, W] dlatents2 = self.fc(dlatents2) # 2D blocks for i in range(len(self.BottleNeck_2d)): x, mask_2d = self.BottleNeck_2d[i](x, dlatents, dlatents2) # => ( [N,32*D,H,W], [N,1,H,W] ) if return_mask: mask_list.append(mask_2d) # reshape back => [N, 32, D, H, W] x = x.view(bs, c, d, h, w) # => [N, 32, D, H, W] # 4) 末端普通3D残差 (无mask) x = self.resblocks_3d(x) # => [N, 32, D, H, W] if return_mask: # 返回 (特征, [mask_2d_1, mask_2d_2, ...]) return x, mask_list else: # 只返回特征 return x # if __name__ == "__main__": # # 假设 x 大小: N=2, C=32, D=4, H=64, W=64 # # x = torch.randn(2, 32, 16, 64, 64) # # dlatents = torch.randn(2, 512) # [N, latent_dim=512] # net = transfer_model(latent_dim=512, n_blocks=4) # # out = net(x, dlatents) # # print("out.shape:", out.shape) # total_params = sum(p.numel() for p in net.parameters()) # print("total parameters:", total_params) # # 预期 [2, 32, 4, 64, 64] # class G3d(nn.Module): # def __init__(self): # super(G3d, self).__init__() # # 下采样路径 # self.downsampling = nn.Sequential( # ResBlock3D_stage3(32, 64), # [B, 32, 16, 64, 64] -> [B, 64, 16, 64, 64] # nn.AvgPool3d(kernel_size=2, stride=2), # -> [B, 64, 8, 32, 32] # ResBlock3D_stage3(64, 128), # -> [B, 128, 8, 32, 32] # nn.AvgPool3d(kernel_size=2, stride=2), # -> [B, 128, 4, 16, 16] # ResBlock3D_stage3(128, 256), # -> [B, 256, 4, 16, 16] # nn.AvgPool3d(kernel_size=2, stride=2), # -> [B, 256, 2, 8, 8] # ResBlock3D_stage3(256, 512), # -> [B, 512, 2, 8, 8] # ) # # 上采样路径 # self.upsampling = nn.Sequential( # ResBlock3D_stage3(512, 256), # -> [B, 256, 2, 8, 8] # nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True), # -> [B, 256, 4, 16, 16] # ResBlock3D_stage3(256, 128), # -> [B, 128, 4, 16, 16] # nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True), # -> [B, 128, 8, 32, 32] # ResBlock3D_stage3(128, 64), # -> [B, 64, 8, 32, 32] # nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True), # -> [B, 64, 16, 64, 64] # ResBlock3D_stage3(64, 32), # -> [B, 32, 16, 64, 64] # ) # # 最终输出层，将通道数恢复为32 # self.final_conv = nn.Sequential( # nn.Conv3d(32, 32, kernel_size=3, padding=1) # 保持最终维度 # ) # def forward(self, x): # """ # 输入: [B, 32, 16, 64, 64] # 输出: [B, 32, 16, 64, 64] # """ # x = self.downsampling(x) # x = self.upsampling(x) # x = self.final_conv(x) # return x # class G3d(nn.Module): # def __init__(self): # super(G3d, self).__init__() # # 简单串联几个3D ResBlock # self.resblocks = nn.Sequential( # ResBlock3D_stage3(32, 32), # 保持通道数不变 # ResBlock3D_stage3(32, 32), # ResBlock3D_stage3(32, 32), # ResBlock3D_stage3(32, 32), # ResBlock3D_stage3(32, 32), # ResBlock3D_stage3(32, 32) # ) # def forward(self, x): # """ # 输入: [B, 32, 16, 64, 64] # 输出: [B, 32, 16, 64, 64] # """ # return self.resblocks(x) class G3d(nn.Module): def __init__(self): super(G3d, self).__init__() # 简单串联几个3D ResBlock self.resblocks1 = nn.Sequential( ResBlock3D_stage3_leak(32, 32), # 保持通道数不变 ResBlock3D_stage3_leak(32, 32), ResBlock3D_stage3_leak(32, 32) ) self.resblocks2 = nn.Sequential( ResBlock2d(32*16, 32*16, kernel_size=(3, 3), padding=(1, 1)), # 保持通道数不变 ResBlock2d(32*16, 32*16, kernel_size=(3, 3), padding=(1, 1)), ResBlock2d(32*16, 32*16, kernel_size=(3, 3), padding=(1, 1)) ) self.resblocks3 = nn.Sequential( ResBlock3D_stage3_leak(32, 32), # 保持通道数不变 ResBlock3D_stage3_leak(32, 32), ResBlock3D_stage3_leak(32, 32) ) # self.third = SameBlock2d(32*16, 32*16, kernel_size=(3, 3), padding=(1, 1), lrelu=True) # self.fourth = nn.Conv2d(in_channels=32*16, out_channels=32*16, kernel_size=1, stride=1) def forward(self, x): """ 输入: [B, 32, 16, 64, 64] 输出: [B, 32, 16, 64, 64] """ x = self.resblocks1(x) bs, c, d, h, w = x.shape x = x.view(bs, c*d, h, w) # => [N, 32*D, H, W] x = self.resblocks2(x) # x = self.fourth(x) x = x.view(bs, c, d, h, w) # => [N, C, D, H, W] x = self.resblocks3(x) return x # class G3d(nn.Module): # def __init__(self): # super(G3d, self).__init__() # # 下采样路径 # self.down1 = ResBlock3D_stage3(32, 64) # [B, 32, 16, 64, 64] -> [B, 64, 16, 64, 64] # self.pool1 = nn.AvgPool3d(kernel_size=2, stride=2) # -> [B, 64, 8, 32, 32] # self.down2 = ResBlock3D_stage3(64, 128) # -> [B, 128, 8, 32, 32] # self.pool2 = nn.AvgPool3d(kernel_size=2, stride=2) # -> [B, 128, 4, 16, 16] # self.down3 = ResBlock3D_stage3(128, 256) # -> [B, 256, 4, 16, 16] # self.pool3 = nn.AvgPool3d(kernel_size=2, stride=2) # -> [B, 256, 2, 8, 8] # self.down4 = ResBlock3D_stage3(256, 512) # -> [B, 512, 2, 8, 8] # # 上采样路径 # self.up1 = ResBlock3D_stage3(512, 256) # [B, 512, 2, 8, 8] -> [B, 256, 2, 8, 8] # self.upsample1 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True) # -> [B, 256, 4, 16, 16] # self.up2 = ResBlock3D_stage3(512, 128) # [B, 512(256+256), 4, 16, 16] -> [B, 128, 4, 16, 16] # self.upsample2 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True) # -> [B, 128, 8, 32, 32] # self.up3 = ResBlock3D_stage3(256, 64) # [B, 256(128+128), 8, 32, 32] -> [B, 64, 8, 32, 32] # self.upsample3 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True) # -> [B, 64, 16, 64, 64] # self.up4 = ResBlock3D_stage3(128, 32) # [B, 128(64+64), 16, 64, 64] -> [B, 32, 16, 64, 64] # # 最终输出层 # self.final_conv = nn.Conv3d(32, 32, kernel_size=3, padding=1) # def forward(self, x): # # 保存初始输入 # x_input = x # [B, 32, 16, 64, 64] # # 下采样路径 # x1 = self.down1(x) # [B, 64, 16, 64, 64] # p1 = self.pool1(x1) # [B, 64, 8, 32, 32] # x2 = self.down2(p1) # [B, 128, 8, 32, 32] # p2 = self.pool2(x2) # [B, 128, 4, 16, 16] # x3 = self.down3(p2) # [B, 256, 4, 16, 16] # p3 = self.pool3(x3) # [B, 256, 2, 8, 8] # x4 = self.down4(p3) # [B, 512, 2, 8, 8] # # 上采样路径 # u1 = self.up1(x4) # [B, 256, 2, 8, 8] # u1 = self.upsample1(u1) # [B, 256, 4, 16, 16] # u1 = torch.cat([u1, x3], dim=1) # [B, 512, 4, 16, 16] # u2 = self.up2(u1) # [B, 128, 4, 16, 16] # u2 = self.upsample2(u2) # [B, 128, 8, 32, 32] # u2 = torch.cat([u2, x2], dim=1) # [B, 256, 8, 32, 32] # u3 = self.up3(u2) # [B, 64, 8, 32, 32] # u3 = self.upsample3(u3) # [B, 64, 16, 64, 64] # u3 = torch.cat([u3, x1], dim=1) # [B, 128, 16, 64, 64] # u4 = self.up4(u3) # [B, 32, 16, 64, 64] # output = self.final_conv(u4) # [B, 32, 16, 64, 64] # return output # ---------------- 测试运行 ---------------- if __name__ == "__main__": # 随机输入测试 N, C, D, H, W = 2, 32, 16, 64, 64 latent_dim = 512 x = torch.randn(N, C, D, H, W) dlatents = torch.randn(N, latent_dim) model = transfer_model(latent_dim=latent_dim, n_blocks=4) # 1) return_mask = False，只返回最终特征 out_no_mask = model(x, dlatents, return_mask=False) print("out_no_mask shape:", out_no_mask.shape) # 2) return_mask = True，同时收集所有 mask out_with_mask, mask_list = model(x, dlatents, return_mask=True) print("out_with_mask shape:", out_with_mask.shape) print(f"mask_list len: {len(mask_list)}") for i, m in enumerate(mask_list): print(f"Mask {i} shape = {m.shape}")