class for BaseTrainer
| 10 | |
| 11 | |
| 12 | class BaseTrainer(object): |
| 13 | """class for BaseTrainer""" |
| 14 | def __init__(self, args): |
| 15 | self.args = args |
| 16 | self.local_group = None |
| 17 | self.logger = None |
| 18 | |
| 19 | def init_ddp_environment(self, gpu, ngpus_per_node): |
| 20 | """ |
| 21 | Args: |
| 22 | gpu: current gpu id |
| 23 | ngpus_per_node: num of process/gpus per node |
| 24 | """ |
| 25 | self.args.gpu = gpu |
| 26 | self.args.ngpus_per_node = ngpus_per_node |
| 27 | self.args.node_rank = self.args.rank |
| 28 | self.args.local_rank = gpu |
| 29 | self.args.local_center = self.args.rank * ngpus_per_node |
| 30 | |
| 31 | torch.cuda.set_device(gpu) |
| 32 | cudnn.benchmark = True |
| 33 | |
| 34 | if self.args.gpu is not None: |
| 35 | print("Use GPU: {} for training".format(self.args.gpu)) |
| 36 | |
| 37 | if self.args.distributed: |
| 38 | if self.args.multiprocessing_distributed: |
| 39 | self.args.rank = self.args.rank * ngpus_per_node + gpu |
| 40 | os.environ['PYTHONWARNINGS'] = 'ignore:semaphore_tracker:UserWarning' |
| 41 | dist.init_process_group( |
| 42 | backend=self.args.dist_backend, init_method=self.args.dist_url, |
| 43 | world_size=self.args.world_size, rank=self.args.rank) |
| 44 | |
| 45 | # setup local group on each node, for ShuffleBN |
| 46 | local_groups = [] |
| 47 | for i in range(0, self.args.world_size // ngpus_per_node): |
| 48 | gp = torch.distributed.new_group( |
| 49 | ranks=list(range(i * ngpus_per_node, (i + 1) * ngpus_per_node)), |
| 50 | backend=self.args.dist_backend) |
| 51 | local_groups.append(gp) |
| 52 | |
| 53 | local_group = local_groups[self.args.rank // ngpus_per_node] |
| 54 | if self.args.local_rank == 0: |
| 55 | print("node_rank:", self.args.node_rank) |
| 56 | print("local_center:", self.args.local_center) |
| 57 | print("local group size:", dist.get_world_size(local_group)) |
| 58 | |
| 59 | self.local_group = local_group |
| 60 | |
| 61 | def init_tensorboard_logger(self): |
| 62 | args = self.args |
| 63 | if args.rank == 0: |
| 64 | self.logger = tb_logger.Logger(logdir=args.tb_folder, flush_secs=2) |
| 65 | |
| 66 | def adjust_learning_rate(self, optimizer, epoch): |
| 67 | args = self.args |
| 68 | lr = args.learning_rate |
| 69 | if args.cosine: |
nothing calls this directly
no outgoing calls
no test coverage detected