models.py
class S3RecModel(nn.Module):
def __init__(self, args):
super(S3RecModel, self).__init__()
# 각종 임베딩 레이어 초기화
# 아이템 임베딩 (상품 ID -> 벡터)
self.item_embeddings = nn.Embedding(
args.item_size, args.hidden_size, padding_idx=0
)
# 속성 임베딩 (상품 속성 -> 벡터)
self.attribute_embeddings = nn.Embedding(
args.attribute_size, args.hidden_size, padding_idx=0
)
# 위치 임베딩 (시퀀스 내 위치 정보)
self.position_embeddings = nn.Embedding(args.max_seq_length, args.hidden_size)
# 인코더와 정규화 레이어들
self.item_encoder = Encoder(args)
self.LayerNorm = LayerNorm(args.hidden_size, eps=1e-12)
self.dropout = nn.Dropout(args.hidden_dropout_prob)
# 4가지 학습 목적을 위한 정규화 레이어들
self.aap_norm = nn.Linear(args.hidden_size, args.hidden_size) # Associated Attribute Prediction
self.mip_norm = nn.Linear(args.hidden_size, args.hidden_size) # Masked Item Prediction
self.map_norm = nn.Linear(args.hidden_size, args.hidden_size) # Masked Attribute Prediction
self.sp_norm = nn.Linear(args.hidden_size, args.hidden_size) # Segment Prediction
# 이진 교차 엔트로피 손실 함수
self.criterion = nn.BCELoss(reduction="none")
self.apply(self.init_weights)
# 예를 들어 영화 ID가 318(쇼생크 탈출)인 경우
item_vector = model.item_embeddings(torch.tensor([318])) # [1, hidden_size] 차원의 벡터
# 위치가 0인 경우 (시퀀스의 첫 번째 위치)
position_vector = model.position_embeddings(torch.tensor([0])) # [1, hidden_size] 차원의 벡터
아래의 주요 학습 태스크로 나뉩니다.
AAP (Associated Attribute Prediction)
def associated_attribute_prediction(self, sequence_output, attribute_embedding):
목적: 아이템의 속성을 예측하는 태스크
예시: 영화 시퀀스가 주어졌을 때 해당 영화의 장르, 감독 등을 맞추는 것
# 예시
sequence = [아바타, 타이타닉, 터미네이터] # 모두 제임스 캐머런 감독 영화
# -> 모델이 "이 시퀀스의 영화들은 제임스 캐머런 감독일 것이다"라고 예측하도록 학습
MIP (Masked Item Prediction)
def masked_item_prediction(self, sequence_output, target_item):
목적: 마스킹된 아이템이 무엇인지 예측
예시: BERT처럼 일부 아이템을 가리고 그게 무엇인지 맞추는 것
# 예시
원본: [아바타, 타이타닉, 터미네이터]
마스킹: [아바타, MASK, 터미네이터]
# -> 모델이 MASK 위치에 "타이타닉"이 올 것이라고 예측하도록 학습
MAP (Masked Attribute Prediction)
def masked_attribute_prediction(self, sequence_output, attribute_embedding):
목적: 마스킹된 아이템의 속성 예측
예시: 가려진 영화의 장르나 감독을 예측
시퀀스: [아바타, MASK, 터미네이터]
# -> MASK 위치의 영화는 "SF 장르일 것이다", "액션 영화일 것이다" 등을 예측
SP (Segment Prediction)
def segment_prediction(self, context, segment):
목적: 연속된 시청 세션이 서로 관련있는지 예측
예시: 사용자의 시청 기록이 하나의 일관된 세션인지 확인
세션1: [아바타, 타이타닉, 터미네이터] # SF/액션 영화 세션
세션2: [노팅힐, 러브액츄얼리, 타이타닉] # 로맨스 영화 세션
# -> 이 두 세션이 서로 관련있는 세션인지 아닌지 판단
사전학습 (Pretrain)
모든 태스크를 함께 사용
각 태스크의 손실값을 모두 합해서 전체 모델 학습
# trainer.py에서 받은 손실값들
aap_loss # 아이템의 속성 예측이 얼마나 정확한지
mip_loss # 마스킹된 아이템 예측이 얼마나 정확한지
map_loss # 마스킹된 아이템의 속성 예측이 얼마나 정확한지
sp_loss # 세그먼트 예측이 얼마나 정확한지
# 이들을 합쳐서 전체 손실 계산
total_loss = aap_loss + mip_loss + map_loss + sp_loss
trainer.py의 pretrain 함수에 사용됩니다.
aap_loss, mip_loss, map_loss, sp_loss = self.model.pretrain(
attributes,
masked_item_sequence,
pos_items,
neg_items,
masked_segment_sequence,
pos_segment,
neg_segment,
)
Finetune 메소드
# 예시 입력 데이터
batch_size = 2
seq_length = 5
input_ids = torch.tensor([
[1, 2, 3, 4, 0], # 첫 번째 사용자 시퀀스 (마지막은 패딩)
[5, 6, 7, 0, 0] # 두 번째 사용자 시퀀스 (뒤의 두 개는 패딩)
])
# 1. 어텐션 마스크 생성
attention_mask = (input_ids > 0).long()
print("Attention mask:")
print(attention_mask)
# 출력:
# tensor([[1, 1, 1, 1, 0],
# [1, 1, 1, 0, 0]])
# 2. Subsequent 마스크 생성 (seq_length=5 예시)
subsequent_mask = torch.triu(torch.ones(1, 5, 5), diagonal=1)
subsequent_mask = (subsequent_mask == 0).long()
print("\nSubsequent mask:")
print(subsequent_mask[0])
# 출력:
# tensor([[1, 0, 0, 0, 0],
# [1, 1, 0, 0, 0],
# [1, 1, 1, 0, 0],
# [1, 1, 1, 1, 0],
# [1, 1, 1, 1, 1]])
# 3. 최종 어텐션 마스크 (두 마스크의 조합)
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) * subsequent_mask
init_weights
# 임베딩 레이어 초기화 예시
embedding = nn.Embedding(1000, 64) # vocab_size=1000, hidden_size=64
model.init_weights(embedding)
# -> 가중치가 평균 0, 표준편차 initializer_range인 정규분포로 초기화됨
# 선형 레이어 초기화 예시
linear = nn.Linear(64, 32)
model.init_weights(linear)
# -> 가중치는 정규분포로, bias는 0으로 초기화됨
# LayerNorm 초기화 예시
layernorm = LayerNorm(64)
model.init_weights(layernorm)
# -> weight는 1로, bias는 0으로 초기화됨
# 1. 사전학습 단계
사용자_시청_기록 = [
"어벤져스",
"MASK", # 아이언맨이 가려짐
"토르",
"MASK", # 캡틴아메리카가 가려짐
"블랙팬서"
]
# 모델은 동시에:
1. 각 영화의 속성 예측 (AAP) -> "마블 영화", "액션 영화" 등
2. 가려진 영화 예측 (MIP) -> "아이언맨", "캡틴아메리카"
3. 가려진 영화의 속성 예측 (MAP) -> "슈퍼히어로물"
4. 시청 세션의 일관성 확인 (SP) -> "이것은 마블 영화 시청 세션이다"
# 2. 미세조정 단계
# 실제 추천 시스템에서:
사용자_기록 = ["어벤져스", "아이언맨", "토르"]
다음_추천 = model.predict(사용자_기록) # "캡틴아메리카" 추천