full resnet50 precision(bf16+amp)

training/benchmarks/resnet50/pytorch/dataloaders/dataloader.py

@@ -8,7 +8,7 @@
 class ToFloat16(object):
 
     def __call__(self, tensor):
-        return tensor.to(dtype=torch.float16)
+        return tensor.to(dtype=torch.bfloat16)
 
 
 def build_train_dataset(config):

training/benchmarks/resnet50/pytorch/train/trainer.py

@@ -63,6 +63,8 @@ def train_one_epoch(self, train_dataloader, eval_dataloader):
         device = self.device
         epoch = self.training_state.epoch
         scaler = self.scaler
+        criterion = torch.nn.CrossEntropyLoss()
+
         print("Epoch " + str(epoch + 1))
         if self.config.distributed:
             train_dataloader.batch_sampler.sampler.set_epoch(epoch)
@@ -76,15 +78,14 @@ def train_one_epoch(self, train_dataloader, eval_dataloader):
 
             batch = self.process_batch(batch, device)
 
+            dist_pytorch.barrier(self.config.vendor)
             pure_start_time = time.time()
             optimizer.zero_grad()
 
             images, target = batch
             if scaler is not None:
                 with torch.cuda.amp.autocast(enabled=True):
                     output = model(images)
-
-                    criterion = torch.nn.CrossEntropyLoss()
                     loss = criterion(output, target)
 
                 scaler.scale(loss).backward()
@@ -102,6 +103,7 @@ def train_one_epoch(self, train_dataloader, eval_dataloader):
                 print("Train Step " + str(step) + "/" + str(len(data_loader)) +
                       ", Loss : " + str(float(loss)))
 
+            dist_pytorch.barrier(self.config.vendor)
             self.training_state.purecomputetime += time.time(
             ) - pure_start_time
 

training/benchmarks/resnet50/pytorch/train/trainer_adapter.py

@@ -26,7 +26,7 @@ def model_to_fp16(model: nn.Module) -> nn.Module:
     # To prevent OOM for model sizes that cannot fit in GPU memory in full precision
     if config.fp16:
         main_proc_print(" > use fp16...")
-        model.half()
+        model.to(torch.bfloat16)
     return model
 
 

training/nvidia/resnet50-pytorch/README.md

@@ -45,5 +45,6 @@
 | A100单机8卡（1x8） | fp32 | bs=256,lr=0.8 | 22653    | 5663    | 5866    | 6105   | 73.5% | 28.3/40.0 |
 | A100单机单卡（1x1） | fp32 | bs=256,lr=0.8 |  | 782 | 795    | 799 |  | 27.6/40.0 |
 | A100两机8卡（2x8） | fp32 | bs=256,lr=0.8 |  | 10576 | 11085    | 11874 | | 27.9/40.0 |
-
+| A100单机8卡（1x8） | amp | bs=512,lr=0.2 | 15312 | 7544    | 7901 | 9567 | 72.7% | 28.6/40.0 |
+| A100单机8卡（1x8） | bf16 | bs=512,lr=0.2 | 14082 | 8203 | 8550 | 9818 | 64.0% | 28.6/40.0 |