PyTorch 1.6.0 update with native AMP (#573)

* PyTorch have Automatic Mixed Precision (AMP) Training.

* Fixed the problem of inconsistent code length indentation

* Fixed the problem of inconsistent code length indentation

* Mixed precision training is turned on by default
This commit is contained in:
Liu Changyu authored and GitHub committed 2020-07-31 10:52:45 -07:00
1 parent 48e15be498
commit c020875b17
2 files changed
+38 -46

No files matched your search

+2 -2
View File
@@ -22,7 +22,7 @@ def init_seeds(seed=0):
cudnn.benchmark = True
def select_device(device='', apex=False, batch_size=None):
def select_device(device='', batch_size=None):
# device = 'cpu' or '0' or '0,1,2,3'
cpu_request = device.lower() == 'cpu'
if device and not cpu_request: # if device requested other than 'cpu'
@@ -36,7 +36,7 @@ def select_device(device='', apex=False, batch_size=None):
if ng > 1 and batch_size: # check that batch_size is compatible with device_count
assert batch_size % ng == 0, 'batch-size %g not multiple of GPU count %g' % (batch_size, ng)
x = [torch.cuda.get_device_properties(i) for i in range(ng)]
s = 'Using CUDA ' + ('Apex ' if apex else '') # apex for mixed precision https://github.com/NVIDIA/apex
s = 'Using CUDA '
for i in range(0, ng):
if i == 1:
s = ' ' * len(s)