[논문 읽기] SeFa - Closed-Form Factorization of Latent Semantics in GANs 핵심 코드 분석

 

transformation matrix A에 대해 조사해야한다. 왜냐면 An의 최댓값을 이미지 Edit을 최대화 시킬 수 있기 때문이다. 그래서 A^TA의 고유벡터와 고유값을 찾는다.

gan_type = parse_gan_type(net.decoder)
layers, boundaries, values = factorize_weight(net.decoder, args.layer_idx)

generator의 weight와 layer 인덱스를 인자로 받고, 그 결과로 layers, boundarys, values를 내 뱉는다. 문맥상 layers는 이미지 Edit에 관여된 layer들을 의미하고, boundarys가 고유 벡터, values가 고유값을 의미하는 듯하다.

위 값들은 실제 코드에 어떻게 쓰일까?

################################################ 이 부분으로 감싼 영역만 보면 된다.

sem_id는 semantic_id의 약자로 layer 18개 중 특정 고유 벡터를 고르기 위해 설정된 값인듯 하다.

이를 토대로하여 모든 latent_code에 수정을 원하는 layer들의 값에 고유벡터에 임의의 상수 distance를 곱한 값을 더해준다. 이 부분이 Image Editing을 극대화시키는 부분이다. 

나머지는 이 코드를 다시 generator에 넣고  이미지를 생성하는 부분.

pil_images = {}
with torch.no_grad():
    for sam_id in tqdm(range(num_sam), desc='Sample ', leave=False):
        pil_images[sam_id] = {}
        code = codes[sam_id:sam_id + 1]
        for sem_id in tqdm(range(num_sem), desc='Semantic ', leave=False):
            pil_images[sam_id][sem_id] = []
            ################################################
            boundary = boundaries[sem_id:sem_id + 1]
            ################################################
            for col_id, d in enumerate(distances, start=1):
                temp_code = code.copy()
                if gan_type == 'pggan':
                    temp_code += boundary * d
                    image = generator(to_tensor(temp_code))['image']
                elif gan_type in ['stylegan', 'stylegan2']:
                    temp_code[:, layers, :] += boundary * d
                    image = generator.synthesis(to_tensor(temp_code))['image']
                elif gan_type in ['stylegan2_psp']:
                    ################################################
                    temp_code[:, layers, :] += boundary * d
                    ################################################
                    images, result_latent = net.decoder([to_tensor(temp_code)],
                                     input_is_latent=True,#True
                                     randomize_noise=True,
                                     return_latents=False)
                    images = net.face_pool(images)

    #             print(images.size())
    #             image = np.array(tensor2im(images))
                    image = np.squeeze([np.array(tensor2im(image)) for image in images], 0)
                    pil_images[sam_id][sem_id].append(Image.fromarray(image))
                vizer_1.set_cell(sem_id * (num_sam + 1) + sam_id + 1, col_id,
                                 image=image)
                vizer_2.set_cell(sam_id * (num_sem + 1) + sem_id + 1, col_id,
                                 image=image)
    #             del images, image
    #             gc.collect()
    #             torch.cuda.empty_cache()

factorize_weight라는 함수를 좀 더 들여다 봐야 layers, boundarys, values의 의미를 이해할 수 있을 듯하다.

 

def factorize_weight(generator, layer_idx='all'):
    """Factorizes the generator weight to get semantics boundaries.

    Args:
        generator: Generator to factorize.
        layer_idx: Indices of layers to interpret, especially for StyleGAN and
            StyleGAN2. (default: `all`)

    Returns:
        A tuple of (layers_to_interpret, semantic_boundaries, eigen_values).

    Raises:
        ValueError: If the generator type is not supported.
    """
    # Get GAN type.
    gan_type = parse_gan_type(generator)

    # Get layers.
    if gan_type == 'pggan':
        layers = [0]
    if gan_type in ['stylegan', 'stylegan2']:
        if layer_idx == 'all':
            layers = list(range(generator.num_layers))
        else:
            layers = parse_indices(layer_idx,
                                   min_val=0,
                                   max_val=generator.num_layers - 1)
    elif gan_type in ['stylegan2_psp']:
        if layer_idx == 'all':
            layers = list(range(generator.n_latent))
        else:
            layers = parse_indices(layer_idx,
                                   min_val=0,
                                   max_val=generator.n_latent - 1)

    # dhkim add
    weights = []
    for idx in layers:
        layer_name = f'layer{idx}'
        if gan_type == 'stylegan2' and idx == generator.num_layers - 1:
            layer_name = f'output{idx // 2}'
        if gan_type == 'pggan':
            weight = generator.__getattr__(layer_name).weight
            weight = weight.flip(2, 3).permute(1, 0, 2, 3).flatten(1)
        if gan_type in ['stylegan', 'stylegan2']:
            weight = generator.synthesis.__getattr__(layer_name).style.weight.T
        elif gan_type in ['stylegan2_psp']:
            if idx < 1:
                weight = generator.conv1.conv.modulation.weight.T
            elif idx < layers[-1]:
                weight = generator.convs[idx-1].conv.modulation.weight.T
            else:
                weight = generator.convs[idx-2].conv.modulation.weight.T

        weights.append(weight.cpu().detach().numpy())
        
    weight = np.concatenate(weights, axis=1).astype(np.float32)
    weight = weight / np.linalg.norm(weight, axis=0, keepdims=True)
    eigen_values, eigen_vectors = np.linalg.eig(weight.dot(weight.T))

    return layers, eigen_vectors.T, eigen_values

고유벡터를 boundaries라는 value로 return 받는다. A부분에 해당하는 weight를 받는다. layer 단위로 이어 줘야하기에 axis=1로 두고 concatenate를 한다. 그 후 크기를 1로 만들어 준 후 A^T와 A의 내적을 계산한다. 이 값에서 고유값과 고유벡터를 구하는 것이다.

 

---

reference

github.com/genforce/sefa

genforce/sefa