Upload liqi files

0ab2566 verified 1 day ago

3.56 kB

## TODO: MLP model，用来降低维度

import torch.nn as nn
import torch
import warnings
import math
import os 

MLP_NAME = "mlp_model.pt"
MLP_OPTIMIZER_NAME = "mlp_optimizer.pt"
MLP_SCHEDULER_NAME = "mlp_scheduler.pt"

def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor

def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)

class LLM2EmbedProjector(nn.Module):
    def __init__(self, in_dim, out_dim, fp16=False):
        super().__init__()
        self.fp16 = fp16
        self.mlp1 = nn.Sequential(
                nn.Linear(in_dim, out_dim),
            )
            
    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        with torch.amp.autocast(device_type='cuda', enabled=self.fp16):
            x = self.mlp1(x)
            x = nn.functional.normalize(x, dim=-1, p=2)
        return x

    def maybe_load_state_dict_from_path(self, ckpt_dir):
        state_dict = None
        ckpt_file = os.path.join(ckpt_dir, MLP_NAME)
        if os.path.exists(ckpt_file):
            state_dict = torch.load(ckpt_file, map_location='cpu')
        if state_dict:
            self.load_state_dict(state_dict)

class LLM2EmbedAddProjector(nn.Module):
    def __init__(self, in_dim, out_dim, fp16=False):
        super().__init__()
        self.fp16 = fp16
        assert in_dim % out_dim == 0, "in_dim must be divisible by out_dim"
        self.out_dim = out_dim
        self.in_dim = in_dim
        self.n = in_dim // out_dim

    def forward(self, x):
        # x shape: (batch_size, in_dim)
        batch_size = x.size(0)
        x_reshaped = x.view(batch_size, self.out_dim, self.n)  # (batch_size, out_dim, n)
        with torch.amp.autocast(device_type='cuda', enabled=self.fp16):
            output = x_reshaped.sum(dim=2)  # sum over the grouped dimensions
        return output  # shape: (batch_size, out_dim)

    def maybe_load_state_dict_from_path(self, ckpt_dir):
        print("No weights for LLM2EmbedAddProjector")

	## TODO: MLP model，用来降低维度

	import torch.nn as nn
	import torch
	import warnings
	import math
	import os

	MLP_NAME = "mlp_model.pt"
	MLP_OPTIMIZER_NAME = "mlp_optimizer.pt"
	MLP_SCHEDULER_NAME = "mlp_scheduler.pt"

	def _no_grad_trunc_normal_(tensor, mean, std, a, b):
	# Cut & paste from PyTorch official master until it's in a few official releases - RW
	# Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
	def norm_cdf(x):
	# Computes standard normal cumulative distribution function
	return (1. + math.erf(x / math.sqrt(2.))) / 2.

	if (mean < a - 2 * std) or (mean > b + 2 * std):
	warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
	"The distribution of values may be incorrect.",
	stacklevel=2)

	with torch.no_grad():
	# Values are generated by using a truncated uniform distribution and
	# then using the inverse CDF for the normal distribution.
	# Get upper and lower cdf values
	l = norm_cdf((a - mean) / std)
	u = norm_cdf((b - mean) / std)

	# Uniformly fill tensor with values from [l, u], then translate to
	# [2l-1, 2u-1].
	tensor.uniform_(2 * l - 1, 2 * u - 1)

	# Use inverse cdf transform for normal distribution to get truncated
	# standard normal
	tensor.erfinv_()

	# Transform to proper mean, std
	tensor.mul_(std * math.sqrt(2.))
	tensor.add_(mean)

	# Clamp to ensure it's in the proper range
	tensor.clamp_(min=a, max=b)
	return tensor

	def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
	return _no_grad_trunc_normal_(tensor, mean, std, a, b)

	class LLM2EmbedProjector(nn.Module):
	def __init__(self, in_dim, out_dim, fp16=False):
	super().__init__()
	self.fp16 = fp16
	self.mlp1 = nn.Sequential(
	nn.Linear(in_dim, out_dim),
	)

	def _init_weights(self, m):
	if isinstance(m, nn.Linear):
	trunc_normal_(m.weight, std=.02)
	if isinstance(m, nn.Linear) and m.bias is not None:
	nn.init.constant_(m.bias, 0)

	def forward(self, x):
	with torch.amp.autocast(device_type='cuda', enabled=self.fp16):
	x = self.mlp1(x)
	x = nn.functional.normalize(x, dim=-1, p=2)
	return x

	def maybe_load_state_dict_from_path(self, ckpt_dir):
	state_dict = None
	ckpt_file = os.path.join(ckpt_dir, MLP_NAME)
	if os.path.exists(ckpt_file):
	state_dict = torch.load(ckpt_file, map_location='cpu')
	if state_dict:
	self.load_state_dict(state_dict)

	class LLM2EmbedAddProjector(nn.Module):
	def __init__(self, in_dim, out_dim, fp16=False):
	super().__init__()
	self.fp16 = fp16
	assert in_dim % out_dim == 0, "in_dim must be divisible by out_dim"
	self.out_dim = out_dim
	self.in_dim = in_dim
	self.n = in_dim // out_dim

	def forward(self, x):
	# x shape: (batch_size, in_dim)
	batch_size = x.size(0)
	x_reshaped = x.view(batch_size, self.out_dim, self.n) # (batch_size, out_dim, n)
	with torch.amp.autocast(device_type='cuda', enabled=self.fp16):
	output = x_reshaped.sum(dim=2) # sum over the grouped dimensions
	return output # shape: (batch_size, out_dim)

	def maybe_load_state_dict_from_path(self, ckpt_dir):
	print("No weights for LLM2EmbedAddProjector")