def test_get_devices(self): # use cpu only devices = ['cpu'] self.assertListEqual(devices, util.get_devices(devices)) # use gpu cuda_available = torch.cuda.is_available() gpu_count = torch.cuda.device_count() if cuda_available and gpu_count >= 1: self.assertListEqual([0], util.get_devices(['cuda:0'])) self.assertListEqual([0], util.get_devices(['cuda:0', 'cuda:1'])) self.assertListEqual(['cpu'], util.get_devices(['cuda:1'])) with self.assertRaises(AssertionError): util.get_devices(['cuda:1', 'cpu']) if cuda_available and gpu_count >= 2: devices = ['cuda:0', 'cuda:1'] self.assertListEqual([0, 1], util.get_devices(devices))
def __init__(self, hps): """ Glow network :param hps: hyper-parameters for this network :type hps: dict """ super().__init__() self.hps = hps self.flow = FlowModel(in_shape=hps.model.image_shape, hidden_channels=hps.model.hidden_channels, K=hps.model.K, L=hps.model.L, permutation=hps.ablation.flow_permutation, coupling=hps.ablation.flow_coupling, actnorm_scale=hps.model.actnorm_scale, lu_decomposition=hps.ablation.lu_decomposition) if hps.ablation.learn_top: nc = self.flow.output_shapes[-1][1] self.learn_top = module.Conv2dZeros(in_channels=2 * nc, out_channels=2 * nc) if hps.ablation.y_condition: nc = self.flow.output_shapes[-1][1] self.y_emb = module.LinearZeros(hps.dataset.num_classes, nc * 2) self.classifier = module.LinearZeros(nc, hps.dataset.num_classes) num_device = len(util.get_devices(self.hps.device.graph, verbose=False)) assert hps.optim.num_batch_train % num_device == 0 self.register_parameter( 'h_top', nn.Parameter( torch.zeros([ hps.optim.num_batch_train // num_device, self.flow.output_shapes[-1][1] * 2, self.flow.output_shapes[-1][2], self.flow.output_shapes[-1][3] ])))
def build(self, training=True): """ Build network :param training: :type training: :return: :rtype: """ # initialize all variables step = 0 state = None result_subdir = None graph, optimizer, scheduler, criterion_dict = None, None, None, None devices = util.get_devices(self.hps.device.graph) data_device = util.get_devices(self.hps.device.data)[0] # build graph graph = model.Glow(self.hps) graph.to('cpu') # load model if graph is not None: # locate or create result subdir if self.hps.general.warm_start and self.hps.general.resume_run_id != "": result_subdir = util.locate_result_subdir( self.hps.general.result_dir, self.hps.general.resume_run_id) if training and result_subdir is None: result_subdir = util.create_result_subdir( self.hps.general.result_dir, desc=self.hps.profile, profile=self.hps) # load pre-trained model on first device if self.hps.general.warm_start: step_or_model_path = None if os.path.exists(self.hps.general.pre_trained): step_or_model_path = self.hps.general.pre_trained elif self.hps.general.resume_step not in [ '', 'best', 'latest' ]: step_or_model_path = int(self.hps.general.resume_step) if step_or_model_path is not None: state = util.load_model(result_subdir, step_or_model_path, graph, device=devices[0]) if not training and state is None: raise RuntimeError('No pre-trained model for inference') # move graph to devices if 'cpu' in devices: graph = graph.cpu() data_device = 'cpu' else: graph = graph.to(devices[0]) print('[Builder] Use {} for model running and {} for data loading'. format(devices[0], data_device)) # setup optimizer and lr scheduler if training and graph is not None: # get optimizer optimizer_name = self.hps.optim.optimizer.lower() assert optimizer_name in self.optimizer_dict.keys(), \ "Unsupported optimizer: {}".format(optimizer_name) # If you need to move a model to GPU via .cuda(), please do so before constructing optimizers for it. optimizer = self.optimizer_dict[optimizer_name]( graph.parameters(), **self.hps.optim.optimizer_args) if state is not None: optimizer.load_state_dict(state['optimizer']) # get lr scheduler scheduler_name = self.hps.optim.lr_scheduler.lower() scheduler_args = self.hps.optim.lr_scheduler_args assert scheduler_name in self.lr_scheduler_dict.keys(), \ "Unsupported lr scheduler: {}".format(scheduler_name) if 'base_lr' not in scheduler_args: scheduler_args['base_lr'] = self.hps.optim.optimizer_args['lr'] scheduler = partial(self.lr_scheduler_dict[scheduler_name], **scheduler_args) return { 'step': step, 'graph': graph, 'optimizer': optimizer, 'scheduler': scheduler, 'devices': devices, 'data_device': data_device, 'result_subdir': result_subdir }