class CustomLoss(nn.Module): # 注意繼承 nn.Module
? ? def __init__(self):
? ? ? ? super(CustomLoss, self).__init__()

? ? def forward(self, x, y):
? ? ? ? # .....這里寫x與y的處理邏輯，即loss的計(jì)算方法
? ? ? ? return loss # 注意最后只能返回Tensor值，且?guī)荻?，?loss.requires_grad == True

示例代碼：

以一個(gè)pytorch求解線性回歸的代碼為例：

import torch
import torch.nn as nn
import numpy as np
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def get_x_y():
? ? np.random.seed(0)
? ? x = np.random.randint(0, 50, 300)
? ? y_values = 2 * x + 21
? ? x = np.array(x, dtype=np.float32)
? ? y = np.array(y_values, dtype=np.float32)
? ? x = x.reshape(-1, 1)
? ? y = y.reshape(-1, 1)
? ? return x, y


class LinearRegressionModel(nn.Module):
? ? def __init__(self, input_dim, output_dim):
? ? ? ? super(LinearRegressionModel, self).__init__()
? ? ? ? self.linear = nn.Linear(input_dim, output_dim) ?# 輸入的個(gè)數(shù)，輸出的個(gè)數(shù)

? ? def forward(self, x):
? ? ? ? out = self.linear(x)
? ? ? ? return out


if __name__ == '__main__':
? ? input_dim = 1
? ? output_dim = 1
? ? x_train, y_train = get_x_y()

? ? model = LinearRegressionModel(input_dim, output_dim)
? ? epochs = 1000 ?# 迭代次數(shù)
? ? optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
? ? model_loss = nn.MSELoss() # 使用MSE作為loss
? ? # 開始訓(xùn)練模型
? ? for epoch in range(epochs):
? ? ? ? epoch += 1
? ? ? ? # 注意轉(zhuǎn)行成tensor
? ? ? ? inputs = torch.from_numpy(x_train)
? ? ? ? labels = torch.from_numpy(y_train)
? ? ? ? # 梯度要清零每一次迭代
? ? ? ? optimizer.zero_grad()
? ? ? ? # 前向傳播
? ? ? ? outputs: torch.Tensor = model(inputs)
? ? ? ? # 計(jì)算損失
? ? ? ? loss = model_loss(outputs, labels)
? ? ? ? # 返向傳播
? ? ? ? loss.backward()
? ? ? ? # 更新權(quán)重參數(shù)
? ? ? ? optimizer.step()
? ? ? ? if epoch % 50 == 0:
? ? ? ? ? ? print('epoch {}, loss {}'.format(epoch, loss.item()))

步驟1：添加自定義的類

我們就用自定義的寫法來寫與MSE相同的效果，MSE計(jì)算公式如下：

添加一個(gè)類：

class CustomLoss(nn.Module):
? ? def __init__(self):
? ? ? ? super(CustomLoss, self).__init__()
? ? ? ? self.mse_loss = nn.MSELoss()

? ? def forward(self, x, y):
? ? ? ? mse_loss = torch.mean(torch.pow((x - y), 2)) # x與y相減后平方，求均值即為MSE
? ? ? ? return mse_loss

步驟2：修改使用的loss函數(shù)

只需要把原始代碼中的：

model_loss = nn.MSELoss() # 使用MSE作為loss

改為：

model_loss = CustomLoss() ?# 自定義loss

即可

完整代碼：

import torch
import torch.nn as nn
import numpy as np
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def get_x_y():
? ? np.random.seed(0)
? ? x = np.random.randint(0, 50, 300)
? ? y_values = 2 * x + 21
? ? x = np.array(x, dtype=np.float32)
? ? y = np.array(y_values, dtype=np.float32)
? ? x = x.reshape(-1, 1)
? ? y = y.reshape(-1, 1)
? ? return x, y


class LinearRegressionModel(nn.Module):
? ? def __init__(self, input_dim, output_dim):
? ? ? ? super(LinearRegressionModel, self).__init__()
? ? ? ? self.linear = nn.Linear(input_dim, output_dim) ?# 輸入的個(gè)數(shù)，輸出的個(gè)數(shù)

? ? def forward(self, x):
? ? ? ? out = self.linear(x)
? ? ? ? return out


class CustomLoss(nn.Module):
? ? def __init__(self):
? ? ? ? super(CustomLoss, self).__init__()
? ? ? ? self.mse_loss = nn.MSELoss()

? ? def forward(self, x, y):
? ? ? ? mse_loss = torch.mean(torch.pow((x - y), 2))
? ? ? ? return mse_loss


if __name__ == '__main__':
? ? input_dim = 1
? ? output_dim = 1
? ? x_train, y_train = get_x_y()

? ? model = LinearRegressionModel(input_dim, output_dim)
? ? epochs = 1000 ?# 迭代次數(shù)
? ? optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
? ? # model_loss = nn.MSELoss() # 使用MSE作為loss
? ? model_loss = CustomLoss() ?# 自定義loss
? ? # 開始訓(xùn)練模型
? ? for epoch in range(epochs):
? ? ? ? epoch += 1
? ? ? ? # 注意轉(zhuǎn)行成tensor
? ? ? ? inputs = torch.from_numpy(x_train)
? ? ? ? labels = torch.from_numpy(y_train)
? ? ? ? # 梯度要清零每一次迭代
? ? ? ? optimizer.zero_grad()
? ? ? ? # 前向傳播
? ? ? ? outputs: torch.Tensor = model(inputs)
? ? ? ? # 計(jì)算損失
? ? ? ? loss = model_loss(outputs, labels)
? ? ? ? # 返向傳播
? ? ? ? loss.backward()
? ? ? ? # 更新權(quán)重參數(shù)
? ? ? ? optimizer.step()
? ? ? ? if epoch % 50 == 0:
? ? ? ? ? ? print('epoch {}, loss {}'.format(epoch, loss.item()))

到此這篇關(guān)于pytorch自定義loss損失函數(shù)的文章就介紹到這了,更多相關(guān)pytorch loss損失函數(shù)內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: