diff --git a/Vision/classification/image/resnet50/examples/train_eager_consistent.sh b/Vision/classification/image/resnet50/examples/train_eager_consistent.sh index 2b6aac5fe..3245fb67c 100644 --- a/Vision/classification/image/resnet50/examples/train_eager_consistent.sh +++ b/Vision/classification/image/resnet50/examples/train_eager_consistent.sh @@ -45,3 +45,4 @@ python3 -m oneflow.distributed.launch \ --num-epochs $EPOCH \ --train-batch-size $TRAIN_BATCH_SIZE \ --val-batch-size $VAL_BATCH_SIZE \ + --scale-grad diff --git a/Vision/classification/image/resnet50/train.py b/Vision/classification/image/resnet50/train.py index c1ba49ba4..459a4b421 100644 --- a/Vision/classification/image/resnet50/train.py +++ b/Vision/classification/image/resnet50/train.py @@ -257,6 +257,7 @@ def train_one_epoch(self): else: top1_acc = 0 + loss = loss * self.world_size self.meter_train_iter(loss, top1_acc) self.cur_batch += 1 @@ -269,7 +270,7 @@ def train_eager(self): if loss.is_global and self.scale_grad: # NOTE(zwx): scale init grad with world_size # because global_tensor.mean() include dividor numel * world_size - loss = loss / self.world_size + # loss = loss / self.world_size loss.backward() for param_group in self.optimizer.param_groups: for param in param_group.parameters: