Change amp autocast to torch.autocast(device, ...) (#2585)

Author: mengfei25

Diff for: models/image_recognition/pytorch/common/inference.py

@@ -99,7 +99,7 @@ def inference(model, dataloader, datatype, args):
                 break
         x = x.to(memory_format=torch.channels_last)
         if args.precision == "bf16":
-            with torch.cpu.amp.autocast(), torch.no_grad():
+            with torch.autocast("cpu", ), torch.no_grad():
                 model = torch.jit.trace(model, x, strict=False)
             model = torch.jit.freeze(model)
         else:
@@ -117,7 +117,7 @@ def inference(model, dataloader, datatype, args):
                 else:
                     images = images.to(memory_format=torch.channels_last)
                 if args.ipex and args.precision == "bf16" and not args.jit:
-                    with torch.cpu.amp.autocast():
+                    with torch.autocast("cpu", ):
                         if i == warmup_iters:
                             with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof, record_function("model_inference"):
                                 output = model(images)
@@ -141,7 +141,7 @@ def inference(model, dataloader, datatype, args):
                 if not args.ipex and not args.jit:
                     images = images.to(datatype).to(memory_format=torch.channels_last)
                 if args.ipex and args.precision == "bf16" and not args.jit:
-                    with torch.cpu.amp.autocast():
+                    with torch.autocast("cpu", ):
                         if i == warmup_iters:
                             with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof, record_function("model_inference"):
                                 output = model(images)

Diff for: models/image_recognition/pytorch/common/main.py

@@ -475,11 +475,11 @@ def main_worker(gpu, ngpus_per_node, args):
                     x = torch.randn(args.batch_size, 3, 224, 224).contiguous(memory_format=torch.channels_last)
                     if args.bf16:
                         x = x.to(torch.bfloat16)
-                        with torch.cpu.amp.autocast(dtype=torch.bfloat16), torch.no_grad():
+                        with torch.autocast("cpu", dtype=torch.bfloat16), torch.no_grad():
                             model = torch.jit.trace(model, x).eval()
                     elif args.fp16:
                         x = x.to(torch.half)
-                        with torch.cpu.amp.autocast(dtype=torch.half), torch.no_grad():
+                        with torch.autocast("cpu", dtype=torch.half), torch.no_grad():
                             model = torch.jit.trace(model, x).eval()
                     else:
                         with torch.no_grad():
@@ -522,7 +522,7 @@ def main_worker(gpu, ngpus_per_node, args):
                         print('[Info] Running torch.compile() with default backend')
                         model = torch.compile(converted_model)
             elif args.bf16:
-                with torch.no_grad(), torch.cpu.amp.autocast(dtype=torch.bfloat16):
+                with torch.no_grad(), torch.autocast("cpu", dtype=torch.bfloat16):
                     x = x.to(torch.bfloat16)
                     if args.ipex:
                         print('[Info] Running torch.compile() BFloat16 with IPEX backend')
@@ -531,7 +531,7 @@ def main_worker(gpu, ngpus_per_node, args):
                         print('[Info] Running torch.compile() BFloat16 with default backend')
                         model = torch.compile(model)
             elif args.fp16:
-                with torch.no_grad(), torch.cpu.amp.autocast(dtype=torch.half):
+                with torch.no_grad(), torch.autocast("cpu", dtype=torch.half):
                     x = x.to(torch.half)
                     if args.ipex:
                         print('[Info] Running torch.compile() FPloat16 with IPEX backend')
@@ -547,7 +547,7 @@ def main_worker(gpu, ngpus_per_node, args):
                     else:
                         print('[Info] Running torch.compile() Float32 with default backend')
                         model = torch.compile(model)
-            with torch.no_grad(), torch.cpu.amp.autocast(enabled=args.bf16 or args.fp16, dtype=torch.half if args.fp16 else torch.bfloat16):
+            with torch.no_grad(), torch.autocast("cpu", enabled=args.bf16 or args.fp16, dtype=torch.half if args.fp16 else torch.bfloat16):
                 y = model(x)
                 y = model(x)
         validate(val_loader, model, criterion, args)
@@ -572,7 +572,7 @@ def main_worker(gpu, ngpus_per_node, args):
             model, optimizer = ipex.optimize(model, optimizer=optimizer, dtype=torch.half, fuse_update_step=False)
 
         if args.inductor:
-            with torch.cpu.amp.autocast(enabled=args.bf16 or args.fp16, dtype=torch.half if args.fp16 else torch.bfloat16):
+            with torch.autocast("cpu", enabled=args.bf16 or args.fp16, dtype=torch.half if args.fp16 else torch.bfloat16):
                 if args.ipex:
                     print('[Info] Running training steps torch.compile() with IPEX backend')
                     model = torch.compile(model, backend="ipex")
@@ -647,11 +647,11 @@ def train(train_loader, val_loader, model, criterion, optimizer, lr_scheduler, a
                 target = target.cuda(args.gpu, non_blocking=True)
 
             if args.bf16:
-                with torch.cpu.amp.autocast(dtype=torch.bfloat16):
+                with torch.autocast("cpu", dtype=torch.bfloat16):
                     output = model(images)
                 output = output.to(torch.float32)
             elif args.fp16:
-                with torch.cpu.amp.autocast(dtype=torch.half):
+                with torch.autocast("cpu", dtype=torch.half):
                     output = model(images)
                 output = output.to(torch.float32)
 
@@ -727,10 +727,10 @@ def run_weights_sharing_model(m, tid, args):
         while num_images < steps:
             start_time = time.time()
             if not args.jit and args.bf16:
-                with torch.cpu.amp.autocast(dtype=torch.bfloat16):
+                with torch.autocast("cpu", dtype=torch.bfloat16):
                     y = m(x)
             elif not args.jit and args.fp16:
-                with torch.cpu.amp.autocast(dtype=torch.half):
+                with torch.autocast("cpu", dtype=torch.half):
                     y = m(x)
             else:
                 y = m(x)
@@ -813,10 +813,10 @@ def validate(val_loader, model, criterion, args):
                         if i >= args.warmup_iterations:
                             end = time.time()
                         if not args.jit and args.bf16:
-                            with torch.cpu.amp.autocast(dtype=torch.bfloat16):
+                            with torch.autocast("cpu", dtype=torch.bfloat16):
                                 output = model(images)
                         elif not args.jit and args.fp16:
-                            with torch.cpu.amp.autocast(dtype=torch.half):
+                            with torch.autocast("cpu", dtype=torch.half):
                                 output = model(images)
                         else:
                             output = model(images)
@@ -852,10 +852,10 @@ def validate(val_loader, model, criterion, args):
                             target = target.cuda(args.gpu, non_blocking=True)
 
                     if not args.jit and args.bf16:
-                        with torch.cpu.amp.autocast(dtype=torch.bfloat16):
+                        with torch.autocast("cpu", dtype=torch.bfloat16):
                             output = model(images)
                     elif not args.jit and args.fp16:
-                        with torch.cpu.amp.autocast(dtype=torch.half):
+                        with torch.autocast("cpu", dtype=torch.half):
                             output = model(images)
 
                     else:

Diff for: models/image_recognition/pytorch/common/main_runtime_extension.py

@@ -411,7 +411,7 @@ def main_worker(gpu, ngpus_per_node, args):
                     x = torch.randn(batch_per_stream, 3, 224, 224).contiguous(memory_format=torch.channels_last)
                     if args.bf16:
                         x = x.to(torch.bfloat16)
-                        with torch.cpu.amp.autocast(), torch.no_grad():
+                        with torch.autocast("cpu", ), torch.no_grad():
                             model = torch.jit.trace(model, x).eval()
                     else:
                         with torch.no_grad():
@@ -492,7 +492,7 @@ def train(train_loader, model, criterion, optimizer, epoch, args):
         # compute output
 
         if args.bf16:
-            with torch.cpu.amp.autocast():
+            with torch.autocast("cpu", ):
                 output = model(images)
             output = output.to(torch.float32)
         else:
@@ -536,7 +536,7 @@ def run_weights_sharing_model(m, tid, args):
         while num_images < steps:
             start_time = time.time()
             if not args.jit and args.bf16:
-                with torch.cpu.amp.autocast():
+                with torch.autocast("cpu", ):
                     y = m(x)
             else:
                 y = m(x)
@@ -600,7 +600,7 @@ def validate(val_loader, model, criterion, args):
                         if i >= args.warmup_iterations:
                             end = time.time()
                         if not args.jit and args.bf16:
-                            with torch.cpu.amp.autocast():
+                            with torch.autocast("cpu", ):
                                 output = model(images)
                         else:
                             output = model(images)
@@ -632,7 +632,7 @@ def validate(val_loader, model, criterion, args):
                     if args.bf16:
                         images = images.to(torch.bfloat16)
                     if not args.jit and args.bf16:
-                        with torch.cpu.amp.autocast():
+                        with torch.autocast("cpu", ):
                             output = model(images)
                     else:
                         output = model(images)

Diff for: models/image_recognition/pytorch/common/train.py

@@ -328,11 +328,11 @@ def main_worker(args):
 
             # Forward pass
             if args.ipex and args.bf16:
-                with torch.cpu.amp.autocast():
+                with torch.autocast("cpu", ):
                     output = model(images)
                 output = output.to(torch.float32)
             elif args.ipex and args.fp16:
-                with torch.cpu.amp.autocast(dtype=torch.half):
+                with torch.autocast("cpu", dtype=torch.half):
                     output = model(images)
                 output = output.to(torch.float32)
             else:
@@ -472,12 +472,12 @@ def validate(val_loader, model, criterion, epoch, args):
 
             if args.ipex and args.bf16:
                 images = images.to(torch.bfloat16)
-                with torch.cpu.amp.autocast():
+                with torch.autocast("cpu", ):
                     output = model(images)
                 output = output.to(torch.float32)
             if args.ipex and args.fp16:
                 images = images.to(torch.half)
-                with torch.cpu.amp.autocast(dtype=torch.half):
+                with torch.autocast("cpu", dtype=torch.half):
                     output = model(images)
                 output = output.to(torch.float32)
             else: 

Diff for: models_v2/pytorch/3d_unet/inference/gpu/predict.py

@@ -173,7 +173,7 @@ def inference_config(model):
             io_utils.write_info('Using JIT trace')
             with torch.inference_mode():
                 if args.xpu and args.ipex:
-                    with torch.xpu.amp.autocast(enabled=use_autocast, dtype=autocast_dtype, cache_enabled=False):
+                    with torch.autocast("xpu", enabled=use_autocast, dtype=autocast_dtype, cache_enabled=False):
                         model = torch.jit.trace(model, trace_input)
                 elif args.gpu or args.xpu:
                     with torch.autocast(enabled=use_autocast, device_type=get_device_type(), dtype=autocast_dtype, cache_enabled=False):
@@ -226,7 +226,7 @@ def do_warmup(model, ds):
                     outputs += [model(images)]
             else:
                 if args.ipex:
-                    with torch.xpu.amp.autocast(enabled=use_autocast, dtype=autocast_dtype, cache_enabled=True):
+                    with torch.autocast("xpu", enabled=use_autocast, dtype=autocast_dtype, cache_enabled=True):
                         for batch_repeat_index in range(min([args.batch_streaming, args.warm_up - len(outputs)])):
                             outputs += [model(images)]
                 else:
@@ -316,7 +316,7 @@ def do_perf_benchmarking(model, ds, gt_data):
                                 torch.xpu.synchronize(args.device)
                                 statistics_utils.accuracy(args, outputs[0], target, overall, whole, core, enhancing, gt_data)
                     else:
-                        with torch.xpu.amp.autocast(enabled=use_autocast, dtype=autocast_dtype, cache_enabled=False):
+                        with torch.autocast("xpu", enabled=use_autocast, dtype=autocast_dtype, cache_enabled=False):
                             start_time = time.time()
                             # compute output
                             for batch_repeat_index in range(args.batch_streaming):

Diff for: models_v2/pytorch/LCM/inference/cpu/inference.py

@@ -71,7 +71,7 @@ def run_weights_sharing_model(pipe, tid, args):
         # run model
         start = time.time()
         if args.precision == "bf16" or args.precision == "fp16" or args.precision == "int8-bf16":
-            with torch.cpu.amp.autocast(dtype=args.dtype), torch.no_grad():
+            with torch.autocast("cpu", dtype=args.dtype), torch.no_grad():
                 output = pipe(args.prompt, generator=torch.manual_seed(args.seed)).images
         else:
             with torch.no_grad():
@@ -180,7 +180,7 @@ def main():
                     if args.precision == "int8-fp32":
                         pipe.unet = ipex.quantization.convert(pipe.unet)
                     else:
-                        with torch.cpu.amp.autocast():
+                        with torch.autocast("cpu", ):
                             pipe.unet = ipex.quantization.convert(pipe.unet)
                     pipe.vae = ipex.optimize(pipe.vae.eval(), dtype=args.dtype, inplace=True)
                     print("running int8 evalation step\n")
@@ -231,14 +231,14 @@ def main():
         print("JIT trace ...")
         # from utils_vis import make_dot, draw
         if args.precision == "bf16" or args.precision == "fp16":
-            with torch.cpu.amp.autocast(dtype=args.dtype), torch.no_grad():
+            with torch.autocast("cpu", dtype=args.dtype), torch.no_grad():
                 pipe.traced_unet = torch.jit.trace(pipe.unet, input, strict=False)
                 pipe.traced_unet = torch.jit.freeze(pipe.traced_unet)
                 pipe.traced_unet(*input)
                 pipe.traced_unet(*input)
                 # print(pipe.traced_unet.graph_for(input))
         elif args.precision == "int8-bf16":
-            with torch.cpu.amp.autocast(), torch.no_grad():
+            with torch.autocast("cpu", ), torch.no_grad():
                 pipe.traced_unet = torch.jit.trace(pipe.unet, input, strict=False)
                 pipe.traced_unet = torch.jit.freeze(pipe.traced_unet)
                 pipe.traced_unet(*input)
@@ -289,14 +289,14 @@ def main():
                 pipe.text_encoder = torch.compile(pipe.text_encoder)
                 pipe.vae.decode = torch.compile(pipe.vae.decode)
         elif args.precision == "bf16":
-            with torch.cpu.amp.autocast(), torch.no_grad():
+            with torch.autocast("cpu", ), torch.no_grad():
                 pipe.unet = torch.compile(pipe.unet)
                 pipe.unet(*input)
                 pipe.unet(*input)
                 pipe.text_encoder = torch.compile(pipe.text_encoder)
                 pipe.vae.decode = torch.compile(pipe.vae.decode)
         elif args.precision == "fp16":
-            with torch.cpu.amp.autocast(dtype=torch.half), torch.no_grad():
+            with torch.autocast("cpu", dtype=torch.half), torch.no_grad():
                 pipe.unet = torch.compile(pipe.unet)
                 pipe.unet(*input)
                 pipe.unet(*input)
@@ -465,7 +465,7 @@ def main():
                     pipe(args.prompt)
                 pipe.traced_unet = convert_pt2e(pipe.traced_unet)
                 torch.ao.quantization.move_exported_model_to_eval(pipe.traced_unet)
-            with torch.cpu.amp.autocast(), torch.no_grad():
+            with torch.autocast("cpu", ), torch.no_grad():
                 pipe.traced_unet = torch.compile(pipe.traced_unet)
                 pipe.traced_unet(*input)
                 pipe.traced_unet(*input)
@@ -493,7 +493,7 @@ def main():
                 # run model
                 start = time.time()
                 if args.precision == "bf16" or args.precision == "fp16" or args.precision == "int8-bf16":
-                    with torch.cpu.amp.autocast(dtype=args.dtype), torch.no_grad():
+                    with torch.autocast("cpu", dtype=args.dtype), torch.no_grad():
                         output = pipe(args.prompt, generator=torch.manual_seed(args.seed)).images
                 else:
                     with torch.no_grad():
@@ -517,7 +517,7 @@ def main():
             real_image = images[0]
             print("prompt: ", prompt)
             if args.precision == "bf16" or args.precision == "fp16" or args.precision == "int8-bf16":
-                with torch.cpu.amp.autocast(dtype=args.dtype), torch.no_grad():
+                with torch.autocast("cpu", dtype=args.dtype), torch.no_grad():
                     output = pipe(prompt, generator=torch.manual_seed(args.seed), output_type="numpy").images
             else:
                 with torch.no_grad():
@@ -546,7 +546,7 @@ def main():
         print("Running profiling ...")
         with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU], record_shapes=True) as p:
             if args.precision == "bf16" or args.precision == "fp16" or args.precision == "int8-bf16":
-                with torch.cpu.amp.autocast(dtype=args.dtype), torch.no_grad():
+                with torch.autocast("cpu", dtype=args.dtype), torch.no_grad():
                     pipe(args.prompt, generator=torch.manual_seed(args.seed)).images
             else:
                 with torch.no_grad():

Diff for: models_v2/pytorch/bert_large/inference/gpu/run_squad.py

@@ -151,7 +151,7 @@ def load_jit_model(model, inputs, dtype, device, jit_trace_path, use_jit_cache):
         in_1 = torch.unsqueeze(inputs["input_ids"][0].clone(), 0)
         in_2 = torch.unsqueeze(inputs["token_type_ids"][0].clone(), 0)
         in_3 = torch.unsqueeze(inputs["attention_mask"][0].clone(), 0)
-        with torch.xpu.amp.autocast(enabled=True, dtype=dtype, cache_enabled=False):
+        with torch.autocast("xpu", enabled=True, dtype=dtype, cache_enabled=False):
             jit_model = torch.jit.trace(model,
                 (in_1.to(device),
                 in_2.to(device),
@@ -396,11 +396,11 @@ def train(args, train_dataset, model, tokenizer):
 
         # autocast context
         if args.device_choice == 'cuda':
-            autocast_context = torch.cuda.amp.autocast(enabled=use_autocast, dtype=autocast_dtype)
+            autocast_context = torch.autocast("cuda", enabled=use_autocast, dtype=autocast_dtype)
         elif args.device_choice == 'xpu':
-            autocast_context = torch.xpu.amp.autocast(enabled=use_autocast, dtype=autocast_dtype)
+            autocast_context = torch.autocast("xpu", enabled=use_autocast, dtype=autocast_dtype)
         else:
-            autocast_context = torch.cpu.amp.autocast(enabled=use_autocast, dtype=autocast_dtype)
+            autocast_context = torch.autocast("cpu", enabled=use_autocast, dtype=autocast_dtype)
 
         import contextlib
         profile_context = contextlib.nullcontext()