forked from ccf-ai-infra/GPUCodeForces
Delete S1/12/.ipynb_checkpoints/tripletmarginloss_torch-checkpoint.py
This commit is contained in:
parent
5a23d186c1
commit
2dc79032ff
|
|
@ -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]
|
||||
Loading…
Reference in New Issue