I have a problem with this GitHub project: https://github.com/researchmm/TTSR
If I use it on one GPU only everything runs smoothly. Once I turn on the second GPU and use torch.nn.DataParallel , this results in "Missing key(s) in state_dict":
[2021-08-03 09:01:00,829] - [trainer.py file line:70] - INFO: Current epoch learning rate: 1.000000e-04
Traceback (most recent call last):
File "/rwthfs/rz/cluster/home/ps815691/git/TTSR/main.py", line 53, in <module>
t.train(current_epoch=epoch, is_init=False)
File "/rwthfs/rz/cluster/home/ps815691/git/TTSR/trainer.py", line 126, in train
sr_lv1, sr_lv2, sr_lv3 = self.model(sr=sr)
File "/home/ps815691/.local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/ps815691/.local/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py", line 167, in forward
outputs = self.parallel_apply(replicas, inputs, kwargs)
File "/home/ps815691/.local/lib/python3.9/site-packages/torch/nn/parallel/data_parallel.py", line 177, in parallel_apply
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
File "/home/ps815691/.local/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
output.reraise()
File "/home/ps815691/.local/lib/python3.9/site-packages/torch/_utils.py", line 429, in reraise
raise self.exc_type(msg)
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
File "/home/ps815691/.local/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
output = module(*input, **kwargs)
File "/home/ps815691/.local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "/rwthfs/rz/cluster/home/ps815691/git/TTSR/model/TTSR.py", line 32, in forward
self.LTE_copy.load_state_dict(self.LTE.state_dict())#, strict=False)
File "/home/ps815691/.local/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1223, in load_state_dict
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for LTE:
Missing key(s) in state_dict: "slice1.0.weight", "slice1.0.bias", "slice2.2.weight", "slice2.2.bias", "slice2.5.weight", "slice2.5.bias", "slice3.7.weight", "slice3.7.bias", "slice3.10.weight", "slice3.10.bias".
I printed the state_dicts for the "LTE" and "LTE_copy":
LTE GPU1 odict_keys([])
LTE GPU0 odict_keys(['sub_mean.weight', 'sub_mean.bias'])
LTE_Copy GPU1 odict_keys([])
LTE_Copy GPU0 odict_keys(['slice1.0.weight', 'slice1.0.bias', 'slice2.2.weight', 'slice2.2.bias', 'slice2.5.weight', 'slice2.5.bias', 'slice3.7.weight', 'slice3.7.bias', 'slice3.10.weight', 'slice3.10.bias', 'sub_mean.weight', 'sub_mean.bias'])
I do not get why that happens. Let me give you a quick introduction to the code: The code starts in main.py. First, the model gets initialized from model/ttsr.py. This ttsr model is composed of several submodels. One of which is "LTE" & "LTE_copy". Then that model is put into nn.DataParallel and the trainer (trainer.py) is initialized with that model. t.train starts the training
_model = TTSR.TTSR(args).to(device)
_model = nn.DataParallel(_model, list(range(args.num_gpu)))
t = Trainer(args, _logger, _dataloader, _model, _loss_all)
t.train(current_epoch=epoch, is_init=True)
In the train function, after a batch has been fed through the model, the models output is fed back to the model, to get some parts of the loss function (trainer.py line 97). The model then executes this code in ttsr.py:
### used in transferal perceptual loss
self.LTE_copy.load_state_dict(self.LTE.state_dict())
sr_lv1, sr_lv2, sr_lv3 = self.LTE_copy((sr + 1.) / 2.)
return sr_lv1, sr_lv2, sr_lv3
Has anyone a clue why the error message above gets thrown out? It does not appear if I use load_state_dict(...,strict=False), but doesn't this just ignore the underlying problem? There does not seem to be any LTE.state_dict on GPU1's memory for example.