New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

多卡训练的问题 #4

Closed

ymj98 opened this issue Nov 2, 2021 · 1 comment

ymj98 commented Nov 2, 2021

您好，请问一下代码是只支持单卡训练么？可以进行多卡训练么

Collaborator

cjy-4 commented Jan 2, 2023 •

edited

Loading

We set the batch size to 1 in our experiments and training on single GPU.
Basically, we have implemented some code here to support multiple GPUs:

BargainNet-Image-Harmonization/models/networks.py

Lines 112 to 116 in faddca2

    
           if len(gpu_ids) > 0: 
        
               assert(torch.cuda.is_available()) 
        
               net.to(gpu_ids[0]) 
        
               net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs 
        
           init_weights(net, init_type, init_gain=init_gain)

You can set --gpu_ids 0,1,... to work on multiple GPUs.

ustcnewly closed this as completed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment