fuse_trainA.py

import argparse
import numpy as np
import os
import json 

import torch
import torch.multiprocessing as mp
import importlib

from option.fuse_optionA import args
import os
import torch
import torch.multiprocessing as mp

import sys
import warnings

warnings.filterwarnings("ignore")

os.environ['CUDA_VISIBLE_DEVICES'] = '3,4'


def main_worker(id, ngpus_per_node, args):
    args.local_rank = args.global_rank = id
    if args.distributed:
        torch.cuda.set_device(args.local_rank)
        print(f'using GPU {args.world_size}-{args.global_rank} for training')
        torch.distributed.init_process_group(
            backend='nccl', init_method=args.init_method,
            world_size=args.world_size, rank=args.global_rank,
            group_name='mtorch')

    args.save_dir = os.path.join(args.save_dir, f'{args.model_name}_{args.file_name}_{args.image_size}')
    
    if (not args.distributed) or args.global_rank == 0:
        os.makedirs(args.save_dir, exist_ok=True)
        with open(os.path.join(args.save_dir, 'config.txt'), 'a') as f:
            for key, val in vars(args).items(): 
                f.write(f'{key}: {val}\n')
        print(f'[**] create folder {args.save_dir}')

    # trainer = Trainer(args)
    trainer = importlib.import_module('trainer.'+args.trainer_name).Trainer(args)
    trainer.train()


if __name__ == "__main__":
    
    torch.manual_seed(args.seed)
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = True

    # setup distributed parallel training environments
    ngpus_per_node = torch.cuda.device_count()
    print('ngpus_per_node is',ngpus_per_node)
    # ngpus_per_node = 2
    # print('ngpus_per_node is',ngpus_per_node)

    if ngpus_per_node > 1:
        args.world_size = ngpus_per_node
        args.init_method = f'tcp://127.0.0.1:{args.port}'
        args.distributed = True
        mp.spawn(main_worker, nprocs=ngpus_per_node,
                 args=(ngpus_per_node, args))
    else:
        args.world_size = 1
        args.distributed = False
        main_worker(0, 1, args)