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
}