戒酒的李白
60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from transformers.models.gpt2.modeling_gpt2 import GPT2Block
|
|
from adapter import AdapterLayer
|
|
|
|
class GPT2BlockWithAdapter(nn.Module):
|
|
"""
|
|
带Adapter的GPT2Block层
|
|
在原始GPT2Block的基础上添加Adapter层实现参数高效微调
|
|
"""
|
|
def __init__(self, config):
|
|
super(GPT2BlockWithAdapter, self).__init__()
|
|
# 创建标准的GPT2Block
|
|
self.original_block = GPT2Block(config)
|
|
|
|
# 添加Adapter层
|
|
adapter_size = 64 # Adapter的隐藏层大小
|
|
self.adapter = AdapterLayer(config.hidden_size, adapter_size)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
layer_past=None,
|
|
attention_mask=None,
|
|
head_mask=None,
|
|
encoder_hidden_states=None,
|
|
encoder_attention_mask=None,
|
|
use_cache=False,
|
|
output_attentions=False,
|
|
**kwargs # 使用**kwargs接收所有其他参数
|
|
):
|
|
# 首先通过原始的GPT2Block，只传递它支持的参数
|
|
outputs = self.original_block(
|
|
hidden_states,
|
|
layer_past=layer_past,
|
|
attention_mask=attention_mask,
|
|
head_mask=head_mask,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
encoder_attention_mask=encoder_attention_mask,
|
|
use_cache=use_cache,
|
|
output_attentions=output_attentions,
|
|
)
|
|
|
|
# 原始输出中的第一个元素是隐藏状态
|
|
hidden_states = outputs[0]
|
|
|
|
# 将隐藏状态通过Adapter层
|
|
hidden_states = self.adapter(hidden_states)
|
|
|
|
# 更新输出的隐藏状态
|
|
outputs = (hidden_states,) + outputs[1:]
|
|
|
|
return outputs
|
|
|
|
def load_state_dict(self, state_dict, strict=True):
|
|
"""
|
|
自定义加载参数方法，用于从原始GPT2Block加载参数
|
|
"""
|
|
# 将所有参数传递给原始Block
|
|
return self.original_block.load_state_dict(state_dict, strict=strict) |