Delete S1/12/.ipynb_checkpoints/tripletmarginloss_torch-checkpoint.py

This commit is contained in:
hli28146 2025-11-11 16:12:13 +08:00
parent 5a23d186c1
commit 2dc79032ff
1 changed files with 0 additions and 41 deletions

View File

@ -1,41 +0,0 @@
# example_torchcode.py
import torch
import torch.nn as nn
import torch.nn.functional as F
class Model(nn.Module):
"""
一个计算Triplet Margin Loss的简单模型
"""
def __init__(self, margin: float = 1.0):
super(Model, self).__init__()
self.margin = margin
self.triplet_margin_loss = torch.nn.TripletMarginLoss(margin=self.margin, reduction='mean')
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
"""
计算三元组损失
使用 reduction='none' 来为批次中的每个样本生成一个损失值以便与CUDA内核进行比较
"""
return self.triplet_margin_loss(anchor, positive, negative)
# 定义标准维度
batch_size = 512
dim = 4096
margin = 1.0
def get_inputs():
"""
为anchor, positive, negative生成三个随机张量
"""
anchor = torch.randn(batch_size, dim)
positive = torch.randn(batch_size, dim)
negative = torch.randn(batch_size, dim)
return [anchor, positive, negative]
def get_init_inputs():
"""
提供模型初始化所需的margin
"""
return [margin]