danielroich / PTI

Official Implementation for "Pivotal Tuning for Latent-based editing of Real Images" (ACM TOG 2022) https://arxiv.org/abs/2106.05744

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Unexpected size of W features in generated_images = self.G.synthesis(w, noise_mode='const', force_fp32=True)

usmancheema89 opened this issue · comments

commented

I am using a self-trained model trained using the StyleGan-ada pytorch repository.
While using use_multi_id_training=True I get a size mismatch error in the forward call of G.synthesis.

The full trace is shown:

Traceback (most recent call last):
File "test.py", line 6, in
run_PTI(use_multi_id_training=True)
File "/home/usman/Documents/Work/PTI/scripts/run_pti.py", line 44, in run_PTI
coach.train()
File "/home/usman/Documents/Work/PTI/training/coaches/multi_id_coach.py", line 58, in train
generated_images = self.forward(w_pivot)
File "/home/usman/Documents/Work/PTI/training/coaches/base_coach.py", line 130, in forward
generated_images = self.G.synthesis(w, noise_mode='const', force_fp32=True)
File "/home/usman/anaconda3/envs/eg3d/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "", line 460, in forward
File "/home/usman/Documents/Work/PTI/torch_utils/misc.py", line 93, in assert_shape
raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}')
AssertionError: Wrong size for dimension 1: got 18, expected 16

commented

Think I figured out the issue.
In w_projector.py change:

del G
return w_opt.repeat([1, 18, 1])

to

G_map_num_ws = G.mapping.num_ws
del G
return w_opt.repeat([1, G_map_num_ws, 1])

Great!
The code correction you suggested is actually the right way to do this. It makes the code agnostic of the generator number of layers.