| |
| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
|
|
| import copy |
| import logging |
| import math |
|
|
| from os.path import join as pjoin |
|
|
| import torch |
| import torch.nn as nn |
| import numpy as np |
|
|
| from torch.nn import BCEWithLogitsLoss,CrossEntropyLoss, Dropout, Softmax, Linear, Conv2d, LayerNorm |
| from torch.nn.modules.utils import _pair |
| from scipy import ndimage |
|
|
| import models.configs as configs |
| from models.attention import Attention |
| from models.embed import Embeddings |
| from models.mlp import Mlp |
|
|
| ATTENTION_Q = "MultiHeadDotProductAttention_1/query" |
| ATTENTION_K = "MultiHeadDotProductAttention_1/key" |
| ATTENTION_V = "MultiHeadDotProductAttention_1/value" |
| ATTENTION_OUT = "MultiHeadDotProductAttention_1/out" |
| FC_0 = "MlpBlock_3/Dense_0" |
| FC_1 = "MlpBlock_3/Dense_1" |
| ATTENTION_NORM = "LayerNorm_0" |
| MLP_NORM = "LayerNorm_2" |
|
|
| class Block(nn.Module): |
| def __init__(self, config, vis, mm=True): |
| super(Block, self).__init__() |
| self.hidden_size = config.hidden_size |
| self.attention_norm = LayerNorm(config.hidden_size, eps=1e-6) |
| self.ffn_norm = LayerNorm(config.hidden_size, eps=1e-6) |
| if mm: |
| self.att_norm_text = LayerNorm(config.hidden_size, eps=1e-6) |
| self.ffn_norm_text = LayerNorm(config.hidden_size, eps=1e-6) |
| self.ffn_text = Mlp(config) |
|
|
| self.ffn = Mlp(config) |
| self.attn = Attention(config, vis, mm) |
|
|
| def forward(self, x, text=None): |
| if text is None: |
| h = x |
| x = self.attention_norm(x) |
| x, text,weights = self.attn(x) |
| |
| x = x + h |
|
|
| h = x |
| x = self.ffn_norm(x) |
| x = self.ffn(x) |
| x = x + h |
| return x |
| else: |
| h = x |
| h_text = text |
| x = self.attention_norm(x) |
| text = self.att_norm_text(text) |
|
|
| x, text, weights_img = self.attn(x, text) |
| |
| x = x + h |
| text = text + h_text |
|
|
| h = x |
| h_text = text |
| x = self.ffn_norm(x) |
| text = self.ffn_norm_text(text) |
| x = self.ffn(x) |
| text = self.ffn_text(text) |
| x = x + h |
| text = text + h_text |
| |
| return x |
|
|
| def load_from(self, weights, n_block): |
| ROOT = f"Transformer/encoderblock_{n_block}" |
| with torch.no_grad(): |
| query_weight = np2th(weights[pjoin(ROOT, ATTENTION_Q, "kernel")]).view(self.hidden_size, self.hidden_size).t() |
| key_weight = np2th(weights[pjoin(ROOT, ATTENTION_K, "kernel")]).view(self.hidden_size, self.hidden_size).t() |
| value_weight = np2th(weights[pjoin(ROOT, ATTENTION_V, "kernel")]).view(self.hidden_size, self.hidden_size).t() |
| out_weight = np2th(weights[pjoin(ROOT, ATTENTION_OUT, "kernel")]).view(self.hidden_size, self.hidden_size).t() |
|
|
| query_bias = np2th(weights[pjoin(ROOT, ATTENTION_Q, "bias")]).view(-1) |
| key_bias = np2th(weights[pjoin(ROOT, ATTENTION_K, "bias")]).view(-1) |
| value_bias = np2th(weights[pjoin(ROOT, ATTENTION_V, "bias")]).view(-1) |
| out_bias = np2th(weights[pjoin(ROOT, ATTENTION_OUT, "bias")]).view(-1) |
|
|
| self.attn.query.weight.copy_(query_weight) |
| self.attn.key.weight.copy_(key_weight) |
| self.attn.value.weight.copy_(value_weight) |
| self.attn.out.weight.copy_(out_weight) |
| self.attn.query.bias.copy_(query_bias) |
| self.attn.key.bias.copy_(key_bias) |
| self.attn.value.bias.copy_(value_bias) |
| self.attn.out.bias.copy_(out_bias) |
|
|
| mlp_weight_0 = np2th(weights[pjoin(ROOT, FC_0, "kernel")]).t() |
| mlp_weight_1 = np2th(weights[pjoin(ROOT, FC_1, "kernel")]).t() |
| mlp_bias_0 = np2th(weights[pjoin(ROOT, FC_0, "bias")]).t() |
| mlp_bias_1 = np2th(weights[pjoin(ROOT, FC_1, "bias")]).t() |
|
|
| self.ffn.fc1.weight.copy_(mlp_weight_0) |
| self.ffn.fc2.weight.copy_(mlp_weight_1) |
| self.ffn.fc1.bias.copy_(mlp_bias_0) |
| self.ffn.fc2.bias.copy_(mlp_bias_1) |
|
|
| self.attention_norm.weight.copy_(np2th(weights[pjoin(ROOT, ATTENTION_NORM, "scale")])) |
| self.attention_norm.bias.copy_(np2th(weights[pjoin(ROOT, ATTENTION_NORM, "bias")])) |
| self.ffn_norm.weight.copy_(np2th(weights[pjoin(ROOT, MLP_NORM, "scale")])) |
| self.ffn_norm.bias.copy_(np2th(weights[pjoin(ROOT, MLP_NORM, "bias")])) |
|
|
|
|
|
|
