CalvinBackup/SimSwapPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix the GPU0 problem

Browse Source
This commit is contained in:
chenxuanhong
2022-02-15 01:40:11 +08:00
parent a148db410c
commit 4a6197a685
10 changed files with 839 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train_scripts/trainer_multi_gpu.py
+1 -3
View File
@@ -5,7 +5,7 @@
# Created Date: Sunday January 9th 2022
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Friday, 11th February 2022 11:18:47 am
# Last Modified: Tuesday, 15th February 2022 12:00:24 am
# Modified By: Chen Xuanhong
# Copyright (c) 2022 Shanghai Jiao Tong University
#############################################################
@@ -215,13 +215,11 @@ def train_loop(
img_std = torch.Tensor([0.229, 0.224, 0.225]).view(3,1,1)
img_mean = torch.Tensor([0.485, 0.456, 0.406]).view(3,1,1)
cudnn_benchmark = True
# Initialize.
device = torch.device('cuda', rank)
np.random.seed(random_seed * num_gpus + rank)
torch.manual_seed(random_seed * num_gpus + rank)
torch.backends.cudnn.benchmark = cudnn_benchmark # Improves training speed.
torch.backends.cuda.matmul.allow_tf32 = False # Improves numerical accuracy.
torch.backends.cudnn.allow_tf32 = False # Improves numerical accuracy.
conv2d_gradfix.enabled = True # Improves training speed.