본문 바로가기
공부/실전문제연구단

[실전문제연구단] GREET 코드 뜯어보기 - 4. train 과정

by 박영귤 2024. 1. 15.

전 게시글들에 모델에 대한 간단한 설명을 적어놓았다. 이번엔 train 과정을 살펴보겠다.

논문에서 소개한 학습 방식의 자료이다.


먼저 Dual-channel Representation Learning Module이다. 이 모듈은 Edge Discriminating Module로부터 Homophiliy한 그래프 View(인접행렬)과 heterophily한 그래프 View(인접행렬)을 받아서 이를 필터에 통과시켜 임베딩을 구한다. 그 후 둘의 차이를 이용해서 Negative Contrastive Loss를 구하여 Back propagation을 진행한다.

def train_cl(cl_model, discriminator, optimizer_cl, features, str_encodings, edges):
    cl_model.train()
    discriminator.eval()

    # adj_1 : low probability
    # adj_2 : high probability
    adj_1, adj_2, weights_lp, _ = discriminator(torch.cat((features, str_encodings), 1), edges)
    # 그래프 뷰를 이용해 새로운 그래프 뷰를 얻어냄.
    features_1, adj_1, features_2, adj_2 = augmentation(features, adj_1, features, adj_2, args, cl_model.training)
    # 두 view 사이의 차이를 이용해 loss 계산
    cl_loss = cl_model(features_1, adj_1, features_2, adj_2)

    optimizer_cl.zero_grad()
    cl_loss.backward()
    optimizer_cl.step()

    return cl_loss.item()

edge discriminator는 Dual-channel Representation Learning Module이 생성한 그래프의 두 가지 View를 이용해 노드임베딩을 계산한다. 그 임베딩을 이용해서, 엣지로 연결된 노도들의 코싸인 유사도와, 랜덤으로 선택한 두 노드의 코싸인 유사도를 계산한다. 이 코싸인 유사도를 이용해서 margin ranking loss를 구한 후, back propagation을 하는 방식이다.

def train_discriminator(cl_model, discriminator, optimizer_disc, features, str_encodings, edges, args):

    cl_model.eval()
    discriminator.train()

    # 판별기를 통해 두 가지 adjacency와 가중치를 계산
    adj_1, adj_2, weights_lp, weights_hp = discriminator(torch.cat((features, str_encodings), 1), edges)

    # 무작위 노드 쌍 생성
    rand_np = generate_random_node_pairs(features.shape[0], edges.shape[1])

    # Positive Sample와 Negative Sample의 레이블 (1: Positive)
    psu_label = torch.ones(edges.shape[1]).to(device)

    # 노드 임베딩 계산
    embedding = cl_model.get_embedding(features, adj_1, adj_2)

    # 엣지로 연결된 노드 쌍의 코사인 유사도를 계산함. 
    edge_emb_sim = F.cosine_similarity(embedding[edges[0]], embedding[edges[1]])

    # 랜덤 노드 쌍의 코사인 유사도를 pivot으로 설정함.
    rnp_emb_sim_lp = F.cosine_similarity(embedding[rand_np[0]], embedding[rand_np[1]])

    # 둘 사이의 ranking loss 계산
    loss_lp = F.margin_ranking_loss(edge_emb_sim, rnp_emb_sim_lp, psu_label, margin=args.margin_hom, reduction='none')
    loss_lp *= torch.relu(weights_lp - 0.5)

    # 랜덤 노드 쌍의 코사인 유사도를 pivot으로 설정함.(lp, hp가 같음. 왜 나눈지는 모르겠음.)
    rnp_emb_sim_hp = F.cosine_similarity(embedding[rand_np[0]], embedding[rand_np[1]])

    # 둘 사이의 ranking loss 계산
    loss_hp = F.margin_ranking_loss(rnp_emb_sim_hp, edge_emb_sim, psu_label, margin=args.margin_het, reduction='none')
    loss_hp *= torch.relu(weights_hp - 0.5)

    # Low Priority와 High Priority Loss의 평균 계산
    rank_loss = (loss_lp.mean() + loss_hp.mean()) / 2

    optimizer_disc.zero_grad()
    rank_loss.backward()
    optimizer_disc.step()

    return rank_loss.item()

 

위 코드는 논문에 작성된 github의 코드이다.

 

의문점

1. rnp_emb_sim_lp, rnp_emb_sim_hp 두 개가 같은 값을 가진다. 왜 둘로 나눠놓은지 모르겠다.

2. loss_lp, loss_hp는 둘 다 margin_ranking_loss함수의 결과값이고, 인풋 임베딩의 순서만 다르게 해서 파라미터도 같은 걸 넣어준다. 이 순서가 어떤 의미인지 모르겠다.(loss값은 다른 값이 나오는 것 확인함)


Reference

논문 - https://arxiv.org/pdf/2211.14065.pdf

github - https://github.com/yixinliu233/GREET 

https://pytorch.org/docs/stable/generated/torch.nn.functional.margin_ranking_loss.html  

https://gombru.github.io/2019/04/03/ranking_loss/