forked from ccf-ai-infra/GPUCodeForces
44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
Variance implementation.
|
|
Computes the variance of input tensors along the feature dimension.
|
|
"""
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Compute variance of input tensor.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor [batch_size, feature_dim]
|
|
|
|
Returns:
|
|
torch.Tensor: Variance values [batch_size]
|
|
"""
|
|
# Compute mean
|
|
mean = torch.mean(x, dim=1, keepdim=True) # [batch_size, 1]
|
|
|
|
# Compute squared differences
|
|
diff = x - mean # [batch_size, feature_dim]
|
|
squared_diff = torch.pow(diff, 2) # [batch_size, feature_dim]
|
|
|
|
# Compute variance
|
|
variance = torch.mean(squared_diff, dim=1) # [batch_size]
|
|
|
|
return variance
|
|
|
|
batch_size = 256
|
|
feature_dim = 1024
|
|
|
|
def get_inputs():
|
|
# Generate input tensor with some variance
|
|
x = torch.randn(batch_size, feature_dim) * 2.0 + 1.0 # mean=1, std=2
|
|
return [x]
|
|
|
|
def get_init_inputs():
|
|
return [] # No special initialization inputs needed
|