ipm-4d-j5的训练编译

import torch from torch import nn from torch.onnx import register_custom_op_symbolic import torch.onnx as onnx N, D_in, D_out = 10, 5, 3 # 自定义一个网络模型 class ONNXSelf(nn.Module): # 初始化函数中对参数进行定义 def __init__(self, in_features, out_features, bias=True): super(ONNXSelf, self).__init__() self.in_features = in_features self.out_features = out_features self.weight = nn.Parameter(torch.Tensor(in_features, out_features)) if bias: self.bias = nn.Parameter(torch.Tensor(in_features)) else: self.register_parameter('bias', None) def forward(self, input): input_ = torch.pow(input, 2) + self.bias y = torch.matmul(input_, self.weight) return y class MyNet(nn.Module): def __init__(self): super(MyNet, self).__init__() self.my_net = ONNXSelf(D_in, D_out) def forward(self, x): x = self.my_net(x) return x model = MyNet() print(model) def main(): x = torch.randn(N, D_in) y = torch.randn(N, D_out) loss_fn = nn.MSELoss(reduction='sum') learning_rate = 1e-4 optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) for i in range(10): y_pred = model(x) loss = loss_fn(y_pred, y) print(f"第{i}个epoch，损失是{loss.item()}") optimizer.zero_grad() loss.backward() optimizer.step() class CustomOperator(torch.autograd.Function): @staticmethod def forward(ctx, input): # 前向传播逻辑 output = torch.add(input , 1) ctx.save_for_backward(input, output) return output register_custom_op_symbolic("custom_operator", CustomOperator, symbolic_fn=None) if __name__ == '__main__': main()