forked from ccf-ai-infra/GPUCodeForces
42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
Weighted Sum implementation.
|
|
Computes the weighted sum of values using corresponding weights.
|
|
"""
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
|
|
def forward(self, values: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Compute weighted sum of values.
|
|
|
|
Args:
|
|
values (torch.Tensor): Input values [batch_size, feature_dim]
|
|
weights (torch.Tensor): Corresponding weights [batch_size, feature_dim]
|
|
|
|
Returns:
|
|
torch.Tensor: Weighted sums [batch_size]
|
|
"""
|
|
# 逐元素乘法
|
|
elementwise_product = values * weights
|
|
|
|
# 求和
|
|
result = torch.sum(elementwise_product, dim=1)
|
|
|
|
return result
|
|
|
|
batch_size = 1024
|
|
feature_dim = 512
|
|
|
|
def get_inputs():
|
|
# Generate values and corresponding weights
|
|
values = torch.randn(batch_size, feature_dim)
|
|
weights = torch.rand(batch_size, feature_dim) # Random weights between 0 and 1
|
|
return [values, weights]
|
|
|
|
def get_init_inputs():
|
|
return [] # No special initialization inputs needed
|