Model/fur_rl/models/retriever_txt.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from Model.CLIP.cn_clip.clip import load_from_name
import cn_clip.clip as clip
from PIL import Image
from torch.autograd import Variable
from fur_rl.models.StageOne import Stage1
from fur_rl.models.transformer import ModelAttn
import copy
import PIL

# v3 带softmax的模型，actor
class Retriever_v3(nn.Module):
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        super(Retriever_v3, self).__init__()
        self.sentence_clip_model2 = sentence_clip_model2
        self.sentence_clip_model_grads = sentence_clip_model2
        self.sentence_clip_preprocess2 = sentence_clip_preprocess2
        self.tokenize = clip.tokenize
        self.Stage1 = Stage1(hid_dim=512)
        self.fc02 = nn.Linear(in_features=512, out_features=1, bias=True)
        self.fc03 = nn.Linear(in_features=256, out_features=64, bias=True)
        self.fc04 = nn.Linear(in_features=64, out_features=1, bias=True)
        self.fc12 = nn.Linear(in_features=512, out_features=1, bias=True)
        self.fc13 = nn.Linear(in_features=256, out_features=64, bias=True)
        self.fc14 = nn.Linear(in_features=64, out_features=1, bias=True)
        # self.fc1 = nn.Linear(in_features=516, out_features=1, bias=True)
        self.fc2 = nn.Linear(in_features=512, out_features=512, bias=True)
        self.bn = nn.BatchNorm1d(num_features=1)
        self.Attn0 = ModelAttn()
        self.Attn1 = ModelAttn()
        self.Attn_his = ModelAttn()
        self.norm = nn.LayerNorm(normalized_shape=512)
        self.norm50 = nn.LayerNorm(normalized_shape=50)

        self.Attn_seq = ModelAttn()
        self.fc_seq0 = nn.Linear(in_features=512, out_features=1, bias=True)
        self.fc_seq1 = nn.Linear(in_features=512, out_features=1, bias=True)

        self.Attn_cross = ModelAttn()
        self.fc_q = nn.Linear(in_features=512, out_features=1, bias=True)
        self.fc_p = nn.Linear(in_features=512, out_features=1, bias=True)

        self.hx_his = None
        self.tanh = nn.Tanh()
        self.logistic = nn.Sigmoid()
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)

        self.device = device
    def img_embed(self, img):
        img_embed = self.sentence_clip_model2.encode_image(img)
        img_embed = img_embed / img_embed.norm(dim=-1, keepdim=True)
        return img_embed

    def txt_embed(self, txt_token):
        txt_embed = self.sentence_clip_model2.encode_text(txt_token)
        txt_embed = txt_embed / txt_embed.norm(dim=-1, keepdim=True)
        return txt_embed

    def txt_embed_grads(self, txt_token):
        txt_embed = self.sentence_clip_model_grads.encode_text(txt_token)
        txt_embed = txt_embed / txt_embed.norm(dim=-1, keepdim=True)
        return txt_embed

    def weight_train(self, img, txt_tokens, txt_matric, actions_matric, ranker, k=50):
        # to do: use img and txt to do the fusion
        p_db, topk_score, topk_action, candidates, candidates_id, candidates_pre, weight_txt = \
            self.get_candidates(txt_matric, actions_matric, k=k, ranker=ranker)
        return weight_txt, p_db


    def get_candidates_greedy(self, txt, actions_matric, k, ranker):
        x = self.Attn_seq(q=txt, k=txt, v=txt)      # [batch, 24, 512]
        x = self.tanh(self.fc_seq0(x))         # [batch, 24, 1]
        x = x.squeeze(dim=2).unsqueeze(dim=1)       # [batch, 1, 24]
        # for baseline
        x = torch.ones(x.shape, device=self.device)     # [batch, 1, 24]
        # for baseline
        weight_txt = x
        x = torch.bmm(x, actions_matric).squeeze(dim=1)  # actions_matric[batch, 24, 2000] x [batch, 2000]
        # normalize
        x = (x - x.mean(dim=-1, keepdim=True)) / x.std(dim=-1, keepdim=True)
        p = self.logistic(x)
        p = torch.clamp(p, min=1e-10, max=1 - 1e-10)        # [batch, 2000]
        topk_score, topk_action = torch.topk(p, k=k, dim=-1)      # [batch, k] [batch, k]
        candidates = torch.zeros((txt.shape[0], k, 512), device=self.device)    # [batch, k, 512]
        candidates_pre = torch.zeros((txt.shape[0], k, 3, 224, 224), device=self.device)    # [batch, k, 3, 224, 224]
        candidates_id = torch.zeros((txt.shape[0], k), device=self.device)    # [batch, k]
        for i in range(k):
            candidates[:, i, :] = ranker.data_emb[topk_action[:, i], :]    # [batch, 1, 512]
            candidates_pre[:, i, :, :, :] = ranker.data_pre[topk_action[:, i], :, :, :]    # [batch, 1, 3, 224, 224]
            candidates_id[:, i] = ranker.data_id[topk_action[:, i]]    # [batch, 1]

        return p, topk_score, topk_action, candidates, candidates_id, candidates_pre, weight_txt

    def get_candidates(self, txt, actions_matric, k, ranker):
        x = self.Attn_seq(q=txt, k=txt, v=txt)      # [batch, 24, 512]
        x = self.tanh(self.fc_seq0(x))         # [batch, 24, 1]
        x = x.squeeze(dim=2).unsqueeze(dim=1)       # [batch, 1, 24]
        # for baseline
        # x = torch.ones(x.shape, device=self.device)     # [batch, 1, 24]
        # for baseline
        weight_txt = x
        x = torch.bmm(x, actions_matric).squeeze(dim=1)  # actions_matric[batch, 24, 2000] x [batch, 2000]
        # normalize
        x = (x - x.mean(dim=-1, keepdim=True)) / x.std(dim=-1, keepdim=True)
        p = self.logistic(x)
        p = torch.clamp(p, min=1e-10, max=1 - 1e-10)        # [batch, 2000]
        topk_score, topk_action = torch.topk(p, k=k, dim=-1)      # [batch, k] [batch, k]
        candidates = torch.zeros((txt.shape[0], k, 512), device=self.device)    # [batch, k, 512]
        candidates_pre = torch.zeros((txt.shape[0], k, 3, 224, 224), device=self.device)    # [batch, k, 3, 224, 224]
        candidates_id = torch.zeros((txt.shape[0], k), device=self.device)    # [batch, k]
        for i in range(k):
            candidates[:, i, :] = ranker.data_emb[topk_action[:, i], :]    # [batch, 1, 512]
            candidates_pre[:, i, :, :, :] = ranker.data_pre[topk_action[:, i], :, :, :]    # [batch, 1, 3, 224, 224]
            candidates_id[:, i] = ranker.data_id[topk_action[:, i]]    # [batch, 1]

        return p, topk_score, topk_action, candidates, candidates_id, candidates_pre, weight_txt


    def action_greedy(self, img, txt_tokens, txt_matric, actions_matric, ranker, k=50):
        # to do: use img and txt to do the fusion
        # txt [batch, turn+2, 52]
        txt_eb = torch.zeros((txt_tokens.shape[0], txt_tokens.shape[1], 512), device=self.device)
        for i in range(txt_tokens.shape[1]):
            txt_eb[:, i, :] = self.txt_embed_grads(txt_tokens[:, i, :].long())
        # txt [batch, 12, 512]
        fusion = torch.cat((txt_eb, img), dim=1)     # [batch, 25, 512]
        p_db, topk_score, topk_action, candidates, candidates_id, candidates_pre, weight_txt = \
            self.get_candidates_greedy(txt_matric, actions_matric, k=k, ranker=ranker)
        x = self.Attn_cross(q=candidates, k=fusion, v=fusion)      # [batch, k, 512]
        # to do: use MLP replace linear
        p = self.softmax(self.fc_p(x)).squeeze(dim=2)       # [batch, k]
        p = torch.clamp(p, min=1e-10, max=1 - 1e-10)
        # to do: use MLP replace linear
        q = self.tanh(self.fc_q(x)).squeeze(dim=2)       # [batch, k]

        p_maxIndex = p.argmax(dim=-1)        # [batch]
        max_candidates = candidates[torch.arange(candidates.shape[0]), p_maxIndex]   # [batch, 512]
        max_candidates_id = candidates_id[torch.arange(candidates_id.shape[0]), p_maxIndex]   # [batch]

        return p, q, max_candidates, max_candidates_id, candidates, candidates_id, candidates_pre, p_db


    def forward(self, img, txt_tokens, txt_matric, actions_matric, ranker, k=50):
        # to do: use img and txt to do the fusion
        # txt [batch, turn+2, 52]
        txt_eb = torch.zeros((txt_tokens.shape[0], txt_tokens.shape[1], 512), device=self.device)
        for i in range(txt_tokens.shape[1]):
            txt_eb[:, i, :] = self.txt_embed_grads(txt_tokens[:, i, :].long())
        # txt [batch, 12, 512]
        fusion = torch.cat((txt_eb, img), dim=1)     # [batch, 25, 512]
        p_db, topk_score, topk_action, candidates, candidates_id, candidates_pre, weight_txt = \
            self.get_candidates(txt_matric, actions_matric, k=k, ranker=ranker)
        x = self.Attn_cross(q=candidates, k=fusion, v=fusion)      # [batch, k, 512]
        # to do: use MLP replace linear
        p = self.softmax(self.fc_p(x)).squeeze(dim=2)       # [batch, k]
        p = torch.clamp(p, min=1e-10, max=1 - 1e-10)
        # to do: use MLP replace linear
        q = self.tanh(self.fc_q(x)).squeeze(dim=2)       # [batch, k]

        p_maxIndex = p.argmax(dim=-1)        # [batch]
        max_candidates = candidates[torch.arange(candidates.shape[0]), p_maxIndex]   # [batch, 512]
        max_candidates_id = candidates_id[torch.arange(candidates_id.shape[0]), p_maxIndex]   # [batch]

        return p, q, max_candidates, max_candidates_id, candidates, candidates_id, candidates_pre, p_db





MEMORY_CAPACITY = 100
N_STATES = 4
GAMMA = 0.95


class DQN_v3():
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device, MULTI_GPU=False, device_ids=None):
        self.actor_net = Retriever_v3(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.num_train = 0
        self.device = device
        self.MULTI_GPU = MULTI_GPU
        self.device_ids = device_ids

        self.memory = []  # initialize memory
        self.neg_memory = []
        self.memory_counter = 0
        self.actor_optimizer = torch.optim.Adam(self.actor_net.parameters(), lr=0.00001)
        self.loss_func1 = nn.MSELoss()
        self.loss_func2 = nn.MSELoss()
        self.batch_mem = {}
        self.is_store = None

    def store_transition(self, g, f, d, a, r, g_, f_, d_, t, txt, txt_, batch_size, net_mem, success_turn, turn):
        if turn == 0:
            for i in range(batch_size):
                self.batch_mem[i] = []
            self.is_store = torch.zeros((batch_size), device=self.device)
        for i in range(batch_size):
            if r[i] <= -2000:
                continue
            else:
                g_tmp = copy.deepcopy(g[i])
                f_tmp = copy.deepcopy(f[i])       # [len, 512]
                d_tmp = copy.deepcopy(d[i])
                a_tmp = copy.deepcopy(a[i])
                reward_temp = copy.deepcopy(r[i])
                g_tmp_ = copy.deepcopy(g_[i])
                f_tmp_ = copy.deepcopy(f_[i])
                d_tmp_ = copy.deepcopy(d_[i])
                t_tmp = copy.deepcopy(t[i])
                txt_tmp = copy.deepcopy(txt[i])
                txt_tmp_ = copy.deepcopy(txt_[i])
                self.batch_mem[i].append((g_tmp, f_tmp, d_tmp, a_tmp, reward_temp,
                                          g_tmp_, f_tmp_, d_tmp_, t_tmp, txt_tmp, txt_tmp_))

                if success_turn[i] > 0 and self.is_store[i] == 0:
                    # print("len", len(self.memory), len(self.batch_mem[i]))
                    self.memory.append(self.batch_mem[i])
                    self.is_store[i] = 1
                    while len(self.memory) > net_mem:
                        self.memory.pop(0)
                    self.memory_counter = len(self.memory)
                elif turn == 9 and self.is_store[i] == 0:
                    self.neg_memory.append(self.batch_mem[i])
                    while len(self.neg_memory) > net_mem:
                        self.neg_memory.pop(0)
                    # print("neg", success_turn[i], reward_temp, len(self.neg_memory))

    def learn(self, batch_size, device=None, ranker=None, k=50, batch_count=0):
        success_turn = 0
        fail_turn = 0
        if len(self.memory) == 0 and len(self.neg_memory) == 0:
            return 0,0
        if len(self.memory) > 0:
            success_turn = 1
            batch_mem = self.memory[0]
            g, f, d, a, r, g_, f_, d_, t, txt, txt_ = zip(*[batch_mem[i] for i in range(len(batch_mem))])
            b_g = torch.stack(g)
            b_f = torch.stack(f)
            b_d = torch.stack(d)
            b_a = torch.stack(a).unsqueeze(1)
            b_reward = torch.stack(r)
            b_r = torch.zeros((len(r), 1), device=self.device)
            for i in range(len(r)):
                if i == 0:
                    b_r[len(r) - i - 1] = r[len(r) - i - 1]
                else:
                    b_r[len(r) - i - 1] = r[len(r) - i - 1] + 0.5 * b_r[len(r) - i]
            b_g_ = torch.stack(g_)
            b_f_ = torch.stack(f_)
            b_d_ = torch.stack(d_)
            b_t = torch.stack(t)
            b_txt = torch.stack(txt)
            b_txt_ = torch.stack(txt_)
            self.memory.pop(0)

        if len(self.neg_memory) > 0:
            fail_turn = 1
            batch_mem = self.neg_memory[0]
            g, f, d, a, r, g_, f_, d_, t, txt, txt_ = zip(*[batch_mem[i] for i in range(len(batch_mem))])
            b_g_neg = torch.stack(g)
            b_f_neg = torch.stack(f)
            b_d_neg = torch.stack(d)
            b_a_neg = torch.stack(a).unsqueeze(1)
            b_reward_neg = torch.stack(r)
            b_r_neg = torch.zeros((len(r), 1), device=self.device)
            for i in range(len(r)):
                if i == 0:
                    b_r_neg[len(r) - i - 1] = r[len(r) - i - 1]
                else:
                    b_r_neg[len(r) - i - 1] = r[len(r) - i - 1] + 0.1 * b_r_neg[len(r) - i]
            b_g_neg_ = torch.stack(g_)
            b_f_neg_ = torch.stack(f_)
            b_d_neg_ = torch.stack(d_)
            b_t_neg = torch.stack(t)
            b_txt_neg = torch.stack(txt)
            b_txt_neg_ = torch.stack(txt_)
            self.neg_memory.pop(0)
        if success_turn > 0 and fail_turn > 0:
            b_g = torch.cat((b_g, b_g_neg), 0)
            b_f = torch.cat((b_f, b_f_neg), 0)
            b_d = torch.cat((b_d, b_d_neg), 0)
            b_a = torch.cat((b_a, b_a_neg), 0)
            b_r = torch.cat((b_r, b_r_neg), 0)
            b_g_ = torch.cat((b_g_, b_g_neg_), 0)
            b_f_ = torch.cat((b_f_, b_f_neg_), 0)
            b_d_ = torch.cat((b_d_, b_d_neg_), 0)
            b_t = torch.cat((b_t, b_t_neg), 0)
            b_reward = torch.cat((b_reward, b_reward_neg), 0)
            b_txt = torch.cat((b_txt, b_txt_neg), 0)
        elif success_turn == 0 and fail_turn > 0:
            b_g = b_g_neg
            b_f = b_f_neg
            b_d = b_d_neg
            b_a = b_a_neg
            b_r = b_r_neg
            b_g_ = b_g_neg_
            b_f_ = b_f_neg_
            b_d_ = b_d_neg_
            b_t = b_t_neg
            b_reward = b_reward_neg
            b_txt = b_txt_neg
        elif success_turn > 0 and fail_turn == 0:
            b_g = b_g
            b_f = b_f
            b_d = b_d
            b_a = b_a
            b_r = b_r
            b_g_ = b_g_
            b_f_ = b_f_
            b_d_ = b_d_
            b_t = b_t
            b_reward = b_reward
            b_txt = b_txt
        ## actor loss
        # img, txt, actions_matric, ranker, k = 50
        p, q, max_candidates, max_candidates_id, candidates, candidates_id, candidates_pre, p_db = \
            self.actor_net(b_g, b_txt, b_f, b_d, ranker, k)
        log_probs = torch.zeros(p.shape[0], 1, device=self.device)
        q_eval = torch.zeros(q.shape[0], 1, device=self.device)
        for i in range(p.shape[0]):
            log_probs[i][0] = p[i][b_a[i][0]].to(device)
            q_eval[i][0] = q[i][b_a[i][0]].to(device)
        # q_eval_temp = torch.zeros((b_a.shape[0], 1)).to(device)
        # for i in range(b_a.shape[0]):
        #     log_probs_temp[i][0] = log_probs[i][b_a[i][0]].to(device)
        #     q_eval_temp[i][0] = log_probs[i][b_a[i][0]].to(device)
        # q_next = self.actor_net(b_g_, b_f_, b_c_, b_hx_, b_r)[1].detach()
        # q_target = b_r + 0.8 * q_next.max(1)[0].view(b_a.shape[0], 1)
        # delta = q_target - q_eval_temp
        # delta = b_r - q_eval_temp
        # print('q_target', q_target[0], 'q_eval', q_eval_temp[0], 'delta', delta[0])
        # print("log_probs_temp", max_score, "b_r", b_r, "log(p)", torch.log(log_probs_temp))
        # actor_loss = - (torch.log(log_probs).t() @ b_r.detach()).float()
        actor_loss = 0
        # critic_loss = torch.mean(delta ** 2).float()
        # critic_loss = torch.mean((b_r/5 - q_eval) ** 2).float()
        critic_loss = 0

        ## supervised_loss
        if success_turn == 1:
            weight, p_db_w = self.actor_net.weight_train(b_g, b_txt, b_f, b_d, ranker, k)
            supervised_loss = torch.zeros(1, device=self.device)
            for i in range(b_reward.shape[0]):
                print(b_reward[i][0], p_db_w[i][b_t[i][0]])
                s_loss_tmp = - b_reward[i][0] * p_db_w[i][b_t[i][0]]
                supervised_loss += s_loss_tmp
            supervised_loss = supervised_loss.float()
        else:
            supervised_loss = 0
        ## total loss
        loss = 0.1 * actor_loss + 10 * critic_loss + 10 * supervised_loss
        if loss != 0:
            self.actor_optimizer.zero_grad()
            print("actor_loss", actor_loss, "critic_loss", critic_loss, "supervised_loss", supervised_loss)
            loss.backward()
            self.actor_optimizer.step()

        return actor_loss, critic_loss