Fix the problem that emanet cannot train with a single card.

Author: haoyuying

configs/emanet/README.md

@@ -11,8 +11,8 @@ Expectation-Maximization Attention Networks for Semantic Segmentation. ICCV 2019
 
 | Model | Backbone | Resolution | Training Iters | mIoU | mIoU (flip) | mIoU (ms+flip) |Links |
 |-|-|-|-|-|-|-|-|
-|EMANet|ResNet50_OS8|1024x512|80000|77.64%|77.98%|78.23%|[model](https://bj.bcebos.com/paddleseg/dygraph/cityscapes/emanet_resnet50_os8_cityscapes_1024x512_80k/model.pdparams) \| [log](https://bj.bcebos.com/paddleseg/dygraph/cityscapes/emanet_resnet50_os8_cityscapes_1024x512_80k/train.log) \| [vdl](https://paddlepaddle.org.cn/paddle/visualdl/service/app?id=3e053a214d60822d6e65445b8614d052)|
-|EMANet|ResNet101_OS8|1024x512|80000|79.41%|79.83%|80.33%|[model](https://bj.bcebos.com/paddleseg/dygraph/cityscapes/emanet_resnet101_os8_cityscapes_1024x512_80k/model.pdparams) \| [log](https://bj.bcebos.com/paddleseg/dygraph/cityscapes/emanet_resnet101_os8_cityscapes_1024x512_80k/train.log) \| [vdl](https://paddlepaddle.org.cn/paddle/visualdl/service/app?id=87be6389cdada711f5c6ada21d9ef6cd)|
+|EMANet|ResNet50_OS8|1024x512|80000|79.05%|79.34%|79.69%|[model](https://bj.bcebos.com/paddleseg/dygraph/cityscapes/emanet_resnet50_os8_cityscapes_1024x512_80k/model.pdparams) \| [log](https://bj.bcebos.com/paddleseg/dygraph/cityscapes/emanet_resnet50_os8_cityscapes_1024x512_80k/train.log) \| [vdl](https://paddlepaddle.org.cn/paddle/visualdl/service/app?id=0a05a0c4cd7d785b9707bdc59f55f585)|
+|EMANet|ResNet101_OS8|1024x512|80000|80.00%|80.23%|80.53%|[model](https://bj.bcebos.com/paddleseg/dygraph/cityscapes/emanet_resnet101_os8_cityscapes_1024x512_80k/model.pdparams) \| [log](https://bj.bcebos.com/paddleseg/dygraph/cityscapes/emanet_resnet101_os8_cityscapes_1024x512_80k/train.log) \| [vdl](https://paddlepaddle.org.cn/paddle/visualdl/service/app?id=ee6926322b8e292ce23ce62ecdaa3439)|
 
 ### Pascal VOC 2012 + Aug
 

paddleseg/models/emanet.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -128,8 +128,6 @@ def __init__(self,
             out_channels=ema_channels,
             kernel_size=3)
         self.ema_mid_conv = nn.Conv2D(ema_channels, ema_channels, kernel_size=1)
-        for param in self.ema_mid_conv.parameters():
-            param.stop_gradient = True
         self.ema_out_conv = layers.ConvBNReLU(
             in_channels=ema_channels, out_channels=ema_channels, kernel_size=1)
         self.bottleneck = layers.ConvBNReLU(
@@ -184,8 +182,8 @@ def __init__(self, c, k, stage_num=3, momentum=0.1):
         tmp_mu = self.create_parameter(
             shape=[1, c, k],
             default_initializer=paddle.nn.initializer.KaimingNormal(k))
-        self.mu = F.normalize(paddle.to_tensor(tmp_mu), axis=1, p=2)
-        self.register_buffer('bases', self.mu)
+        mu = F.normalize(paddle.to_tensor(tmp_mu), axis=1, p=2)
+        self.register_buffer('mu', mu)
 
     def forward(self, x):
         x_shape = paddle.shape(x)
@@ -207,9 +205,11 @@ def forward(self, x):
 
         if self.training:
             mu = paddle.mean(mu, 0, keepdim=True)
-            if paddle.distributed.get_world_size() > 1:
-                paddle.distributed.reduce(
-                    mu / paddle.distributed.get_world_size(), 0)
             mu = F.normalize(mu, axis=1, p=2)
-            self.mu = self.mu * (1 - self.momentum) + mu * self.momentum
+            mu = self.mu * (1 - self.momentum) + mu * self.momentum
+            if paddle.distributed.get_world_size() > 1:
+                mu = paddle.distributed.all_reduce(mu)
+                mu /= paddle.distributed.get_world_size()
+            self.mu = mu
+
         return x