DAST 代码分析
DA部分
输入图片大小：
images.size: torch.Size([1, 3, 512, 1024])
labels.size: torch.Size([1, 512, 1024])
input_size = (w, h) # input_size : <class 'tuple'>: (1024, 512)
input_size_target = (w, h) # <class 'tuple'>: (1024, 512)
分割后的特征图大小：
feat_source: ([1, 2048, 65, 129])
pred_source: ([1, 19, 65, 129])
pred_source = interp(pred_source) 上采样后 pred_source 大小变成: ([1, 19, 512, 1024])
创建网络：
1 model = DeeplabMulti(num_classes=args.num_classes) 2 def DeeplabMulti(num_classes=21): 3 model = ResNetMulti(Bottleneck, [3, 4, 23, 2, 1], num_classes) 4 return model
包含注意力的分割网络：
1 class ResNetMulti(nn.Module): 2 3 def forward(self, x, D, domain): # 源域进来就正常打分, 目标域进来就先加权后打分 4 x = self.conv1(x) 5 x = self.bn1(x) 6 x = self.relu(x) 7 x = self.maxpool(x) 8 x1 = self.layer1(x) 9 x2 = self.layer2(x1) 10 x3 = self.layer3(x2) 11 x4 = self.layer4(x3) # ft或者fs 12 if domain == 'source': # source:x4.size: torch.Size([1, 2048, 65, 129]) out.size: torch.Size([1, 19, 65, 129]) 13 x4_a4 = x4 14 # 目标域 注意力图加权 15 if domain == 'target': # target:x4.size: torch.Size([1, 2048, 65, 129]) out.size: torch.Size([1, 19, 65, 129]) 16 a4 = D[0](x4) #a4 等同于论文中的D（ft） 注意力图 17 a4 = self.tanh(a4) # 防止早期训练时梯度爆炸，tanh激活层作为正则化层 18 a4 = torch.abs(a4) # 绝对值 a4 = |D（ft）| 19 a4_big = a4.expand(x4.size()) # 即a',为了匹配目标域的维度，实现注意力图和目标域按元素相乘 20 x4_a4 = a4_big*x4 + x4 # ft'=ft+ft*a' 21 x5 = self.layer5(x4_a4) 22 out = self.layer6(x5) 23 # print('D[0]',D[0]) 24 # print('domain:', domain) 25 # print('x4.size:', x4.size()) # x4.size: torch.Size([1, 2048, 65, 129]) 26 # print('out.size:', out.size()) # out.size: torch.Size([1, 19, 65, 129]) 27 return x4, out
判别器：（FCDiscriminator输入通道2048，而OutspaceDiscriminator输入通道是19）
model_D = nn.ModuleList([FCDiscriminator(num_classes=num_class_list[i]).train().to( device) if i < 1 else OutspaceDiscriminator(num_classes=num_class_list[i]).train().to(device) for i in range(2)])
class FCDiscriminator(nn.Module):
def __init__(self, num_classes, ndf = 64):
# print('num_classes:', num_classes) num_classes: 2048
super(FCDiscriminator, self).__init__()
self.conv1 = nn.Conv2d(num_classes, num_classes//2, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.Conv2d(num_classes//2, num_classes//4, kernel_size=3, stride=1, padding=1)
self.conv3 = nn.Conv2d(num_classes//4, num_classes//8, kernel_size=3, stride=1, padding=1)
self.classifier = nn.Conv2d(num_classes//8, 1, kernel_size=3, stride=1, padding=1)
self.leaky_relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
#self.up_sample = nn.Upsample(scale_factor=32, mode='bilinear')
#self.sigmoid = nn.Sigmoid()
def forward(self, x):
x = self.conv1(x)
x = self.leaky_relu(x)
x = self.conv2(x)
x = self.leaky_relu(x)
x = self.conv3(x)
x = self.leaky_relu(x)
x = self.classifier(x)
#x = self.up_sample(x)
#x = self.sigmoid(x)
return x
class OutspaceDiscriminator(nn.Module):
def __init__(self, num_classes, ndf = 64):
super(OutspaceDiscriminator, self).__init__()
self.conv1 = nn.Conv2d(num_classes, ndf, kernel_size=4, stride=2, padding=1)
self.conv2 = nn.Conv2d(ndf, ndf*2, kernel_size=4, stride=2, padding=1)
self.conv3 = nn.Conv2d(ndf*2, ndf*4, kernel_size=4, stride=2, padding=1)
self.conv4 = nn.Conv2d(ndf*4, ndf*8, kernel_size=4, stride=2, padding=1)
self.classifier = nn.Conv2d(ndf*8, 1, kernel_size=4, stride=2, padding=1) # 变成通道数为1
self.leaky_relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
#self.up_sample = nn.Upsample(scale_factor=32, mode='bilinear')
#self.sigmoid = nn.Sigmoid()
def forward(self, x):
x = self.conv1(x)
x = self.leaky_relu(x)
x = self.conv2(x)
x = self.leaky_relu(x)
x = self.conv3(x)
x = self.leaky_relu(x)
x = self.conv4(x)
x = self.leaky_relu(x)
x = self.classifier(x)
#x = self.up_sample(x)
#x = self.sigmoid(x)
return x
1 # D[0](x4): 2 # tensor([[[[0.0710, 0.1864, 0.2138, ..., 0.2505, 0.1997, 0.1675], 3 # [0.0946, 0.2139, 0.2130, ..., 0.2979, 0.2266, 0.1543], 4 # [0.1402, 0.2508, 0.2545, ..., 0.3649, 0.3104, 0.1574], 5 # ..., 6 # [0.1940, 0.3481, 0.3824, ..., 0.3082, 0.2303, 0.1237], 7 # [0.1855, 0.2981, 0.3047, ..., 0.2617, 0.1878, 0.0770], 8 # [0.0597, 0.1503, 0.1717, ..., 0.1718, 0.1432, 0.0634]]]], 9 # device='cuda:0', grad_fn= < AddBackward0 >)
1 # D[0]: 2 # FCDiscriminator( 3 # (conv1): Conv2d(2048, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) 4 # (conv2): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) 5 # (conv3): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) 6 # (classifier): Conv2d(256, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) 7 # (leaky_relu): LeakyReLU(negative_slope=0.2, inplace=True) 8 # )
1 # model_D[1]: OutspaceDiscriminator( 2 # (conv1): Conv2d(19, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) 3 # (conv2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) 4 # (conv3): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) 5 # (conv4): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) 6 # (classifier): Conv2d(512, 1, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1)) 7 # (leaky_relu): LeakyReLU(negative_slope=0.2, inplace=True) 8 # )
#######开始训练######
train S
# train with source
feat_source, pred_source = model(images, model_D, 'source')
# ResNet返回两层输出结果(特征图)，x4和out model_D判别器模型，用来对resnet的x4层作输出，得到注意力图
# print('feat_source',feat_source) feat_source=x4 倒数第二层(输出通道数：2048) ； pred_source=out 最后一层（输出通道数：19（类别数）） 如果是目标域图片，那么，out代表加权后的特征图输出。
pred_source = interp(pred_source) # 源域特征图 参与分割损失计算
loss_seg = seg_loss(pred_source, labels)
loss_seg.backward()
# train with target
feat_target, pred_target = model(images, model_D, 'target') # print('feat_target.size, pred_target.size:', feat_target.size(), pred_target.size()) # feat_target.size, pred_target.size: torch.Size([1, 2048, 65, 129]) torch.Size([1, 19, 65, 129]) pred_target = interp_target(pred_target) loss_adv = 0 D_out = model_D[0](feat_target) # 对倒数第二层的T域特征图打分 loss_adv += bce_loss(D_out, torch.FloatTensor(D_out.data.size()).fill_(source_label).to(device)) D_out = model_D[1](F.softmax(pred_target, dim=1)) # 先把最后一层特征图变成概率图，再对概率图打分 # print('model_D[1]:', model_D[1]) loss_adv += bce_loss(D_out, torch.FloatTensor(D_out.data.size()).fill_(source_label).to(device)) loss_adv = loss_adv * 0.01 loss_adv.backward() optimizer.step()
train D
# train with source
loss_D_source = 0 D_out_source = model_D[0](feat_source.detach()) loss_D_source += bce_loss(D_out_source, torch.FloatTensor(D_out_source.data.size()).fill_(source_label).to(device)) D_out_source = model_D[1](F.softmax(pred_source.detach(), dim=1)) loss_D_source += bce_loss(D_out_source, torch.FloatTensor(D_out_source.data.size()).fill_(source_label).to(device)) loss_D_source.backward()
# train with target
loss_D_target = 0 D_out_target = model_D[0](feat_target.detach()) loss_D_target += bce_loss(D_out_target, torch.FloatTensor(D_out_target.data.size()).fill_(target_label).to(device)) D_out_target = model_D[1](F.softmax(pred_target.detach(), dim=1)) loss_D_target += bce_loss(D_out_target, torch.FloatTensor(D_out_target.data.size()).fill_(target_label).to(device)) loss_D_target.backward() optimizer_D.step()
ST部分
images.size: torch.Size([1, 3, 512, 1024])
