用pytorch写一个决策树回归

下面是一个使用 PyTorch 实现的决策树回归的示例代码：

import torch
import numpy as np

class DecisionTreeRegressor:
    def __init__(self, max_depth=3):
        self.max_depth = max_depth
        self.left = None
        self.right = None
        self.split_idx = None
        self.split_val = None
        self.y_val = None
    
    def fit(self, X, y):
        if self.max_depth == 0 or len(y) < 2:
            self.y_val = y.mean()
            return

        best_split_idx, best_split_val = self._best_split(X, y)
        if best_split_idx is None:
            self.y_val = y.mean()
            return
        
        X_left, X_right, y_left, y_right = self._split_data(X, y, best_split_idx, best_split_val)
        
        self.split_idx = best_split_idx
        self.split_val = best_split_val
        self.left = DecisionTreeRegressor(self.max_depth-1)
        self.left.fit(X_left, y_left)
        self.right = DecisionTreeRegressor(self.max_depth-1)
        self.right.fit(X_right, y_right)
        
    def predict(self, X):
        if self.y_val is not None:
            return torch.tensor([self.y_val]*len(X))
        else:
            return torch.where(X[:, self.split_idx] <= self.split_val,
                               self.left.predict(X),
                               self.right.predict(X))
        
    def _best_split(self, X, y):
        best_idx, best_val = None, None
        best_loss = float("inf")
        for i in range(X.shape[1]):
            unique_vals = np.unique(X[:, i])
            for val in unique_vals:
                X_left, X_right, y_left, y_right = self._split_data(X, y, i, val)
                loss = np.mean((y_left - y_left.mean())**2) + np.mean((y_right - y_right.mean())**2)
                if loss < best_loss:
                    best_loss = loss
                    best_idx = i
                    best_val = val
        return best_idx, best_val
    
    def _split_data(self, X, y, idx, val):
        idx_left = X[:, idx] <= val
        idx_right = ~idx_left
        return X[idx_left], X[idx_right], y[idx_left], y[idx_right]

使用方法类似于其他机