代碼功能概述

1. 導(dǎo)入相關(guān)包與設(shè)置環(huán)境變量：
   • 首先導(dǎo)入了如?os、numpy、pandas?等常用的 Python 庫，同時設(shè)置了一些與特定庫（如?XLA_PYTHON_CLIENT_PREALLOCATE?和?JAX_PM AP_USE_TENSORSTORE）相關(guān)的環(huán)境變量，用于優(yōu)化計算等操作。
2. 加載預(yù)訓(xùn)練的 TimesFM 模型：
   • 通過指定相關(guān)超參數(shù)（如后端為?gpu、每核心批處理大小等）以及預(yù)訓(xùn)練模型在 Hugging Face 上的倉庫?id，實例化了?TimesFm?模型對象，用于后續(xù)的評估和微調(diào)等操作。
3. 準(zhǔn)備數(shù)據(jù)集相關(guān)信息并加載數(shù)據(jù)：
   • 定義了一個數(shù)據(jù)集字典?DATA_DICT，包含多個數(shù)據(jù)集（如?ettm1?等）的詳細(xì)信息，包括數(shù)據(jù)文件路徑、時間頻率、劃分邊界等。
   • 根據(jù)選定的數(shù)據(jù)集（示例中初始化為?ettm1），讀取對應(yīng)的數(shù)據(jù)文件為?DataFrame，然后配置?TimeSeriesdata?類的實例來進行數(shù)據(jù)加載、劃分訓(xùn)練集、驗證集和測試集，同時對數(shù)據(jù)進行了一些規(guī)范化等預(yù)處理操作，并生成對應(yīng)的批次數(shù)據(jù)（train_batches、val_batches、test_batches）。
4. 評估預(yù)訓(xùn)練模型在測試集上的 MAE（平均絕對誤差）：
   • 通過迭代測試集批次數(shù)據(jù)，利用預(yù)訓(xùn)練模型進行預(yù)測，計算預(yù)測值和實際值之間的平均絕對誤差，以此來評估模型在當(dāng)前數(shù)據(jù)集上的性能表現(xiàn)。
5. 微調(diào)模型：
   • 導(dǎo)入了一系列用于構(gòu)建和訓(xùn)練模型的?praxis、paxml?相關(guān)的模塊和類，進行了諸如定義學(xué)習(xí)器（包括優(yōu)化器、學(xué)習(xí)率調(diào)度等配置）、構(gòu)建任務(wù)、初始化模型狀態(tài)等操作。
   • 將預(yù)訓(xùn)練模型的參數(shù)設(shè)置為微調(diào)模型的初始權(quán)重，然后通過定義訓(xùn)練步和評估步函數(shù)，在多個?epoch?內(nèi)循環(huán)進行訓(xùn)練和定期評估（利用早停機制，根據(jù)驗證集損失決定是否提前停止訓(xùn)練），在每個訓(xùn)練步中對模型參數(shù)進行更新，每個評估步計算驗證集上的損失，保存最優(yōu)模型狀態(tài)的檢查點。
6. 加載并評估微調(diào)后的模型：
   • 從保存的檢查點中恢復(fù)最優(yōu)的模型狀態(tài)，將其參數(shù)更新到原?TimesFM?模型中，然后再次在測試集上計算平均絕對誤差，以對比微調(diào)前后模型性能的變化情況。

針對車輛銷售數(shù)據(jù)的改寫步驟

1. 數(shù)據(jù)準(zhǔn)備與加載部分（適配車輛銷售數(shù)據(jù)）：

   • 修改數(shù)據(jù)集字典?DATA_DICT：
     • 創(chuàng)建一個新的字典項來對應(yīng)你的車輛銷售數(shù)據(jù)集，例如取名為?vehicle_sales。
     • 填寫對應(yīng)的數(shù)據(jù)文件路徑（假設(shè)你的車輛銷售數(shù)據(jù)存儲在?../datasets/vehicle_sales.csv，則?data_path?設(shè)置為此路徑）。
     • 根據(jù)你的數(shù)據(jù)時間粒度來設(shè)置?freq，比如如果是按天記錄的，就可以設(shè)置為?"D"（代表 Daily），如果是按月記錄的，可設(shè)置為?"M"（代表 Monthly）等。
     • 按照你的訓(xùn)練、驗證、測試集劃分的時間范圍來設(shè)置?boundaries?列表中的值，例如如果前 3 年數(shù)據(jù)作為訓(xùn)練集，第 4 年作為驗證集，第 5 年作為測試集，你需要根據(jù)數(shù)據(jù)點數(shù)量等信息確定對應(yīng)的時間點邊界值填入該列表。
   • 調(diào)整數(shù)據(jù)讀取和數(shù)據(jù)加載配置部分：
     • 在?data_df = pd.read_csv(open(data_path, "r"))?這行代碼中，確認(rèn)數(shù)據(jù)文件格式正確能被?read_csv?方法讀取，如果數(shù)據(jù)有特定的分隔符、編碼等情況，按需調(diào)整參數(shù)（比如添加?sep?參數(shù)指定分隔符、encoding?參數(shù)指定編碼格式等）。
     • 根據(jù)你的車輛銷售數(shù)據(jù)列名，修改?ts_cols、num_cov_cols、cat_cov_cols?的定義。例如，銷售量和銷售價格等數(shù)值型的時間序列列可添加到?ts_cols，車型、經(jīng)銷商這些分類列可以根據(jù)需求分配到?num_cov_cols（如果要進行數(shù)值編碼等處理）或者?cat_cov_cols（作為分類特征）中。同時修改?TimeSeriesdata?實例化時傳入的參數(shù)，確保數(shù)據(jù)能正確劃分和預(yù)處理，例如?datetime_col?設(shè)置為數(shù)據(jù)中代表日期的列名。

2. 模型微調(diào)部分（可能無需大改，但檢查配置合理性）：

   • 確認(rèn)微調(diào)時定義的學(xué)習(xí)器配置（如優(yōu)化器、學(xué)習(xí)率調(diào)度等參數(shù)）是否適合車輛銷售數(shù)據(jù)預(yù)測任務(wù)。你可能需要根據(jù)實際情況調(diào)整學(xué)習(xí)率、訓(xùn)練總步數(shù)等參數(shù)，例如車輛銷售數(shù)據(jù)如果比較復(fù)雜，可能需要適當(dāng)調(diào)小學(xué)習(xí)率、增加訓(xùn)練總步數(shù)等，以保證模型能更好地收斂和學(xué)習(xí)到數(shù)據(jù)中的模式。
   • 檢查?build_learner?函數(shù)中設(shè)置的?bprop_variable_exclusion?參數(shù)是否合理，對于車輛銷售數(shù)據(jù)微調(diào)場景下想要固定或者放開訓(xùn)練的模型層，根據(jù)模型結(jié)構(gòu)和需求進行調(diào)整，確保只訓(xùn)練希望更新參數(shù)的那些部分。

3. 模型評估部分（保持邏輯基本一致）：

   • 在計算微調(diào)前后模型在測試集上的平均絕對誤差（MAE）部分，確保數(shù)據(jù)維度等處理符合車輛銷售數(shù)據(jù)的特點。例如，在預(yù)測結(jié)果和實際結(jié)果對比計算?MAE?時，確認(rèn)預(yù)測的銷售量、銷售價格等和實際值的對應(yīng)關(guān)系和維度對齊正確，特別是如果有多個時間序列維度或者特征維度時，保證?forecasts?和?actuals?的形狀匹配能正確計算誤差。

以下是一個python代碼（假設(shè)你的車輛銷售數(shù)據(jù)?vehicle_sales.csv?有?date（日期）、car_model（車型）、dealer（經(jīng)銷商）、sales_volume（銷售量）、sales_price（銷售價格）這幾列，并且想將車型和經(jīng)銷商作為分類特征，銷售量和銷售價格作為時間序列特征，數(shù)據(jù)按年劃分訓(xùn)練、驗證、測試集，這里簡化假設(shè)前 3 年訓(xùn)練、第 4 年驗證、第 5 年測試，并且時間頻率是按年?"Y"）：

# 以下代碼主要用于基于TimesFM模型對車輛銷售數(shù)據(jù)進行預(yù)訓(xùn)練模型評估、模型微調(diào)以及微調(diào)后模型的評估操作
# 導(dǎo)入相關(guān)包用于微調(diào)操作，同時設(shè)置一些環(huán)境變量來優(yōu)化計算等相關(guān)配置
## Importing relevant packages for finetuning
import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
os.environ['JAX_PMAP_USE_TENSORSTORE'] = 'false'
import timesfm
import gc
import numpy as np
import pandas as pd
from timesfm import patched_decoder
from timesfm import data_loader
from tqdm import tqdm
import dataclasses
import IPython
import IPython.display
import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.rcParams['figure.figsize'] = (8, 6)
mpl.rcParams['axes.grid'] = False# 加載預(yù)訓(xùn)練的TimesFM模型，通過指定相關(guān)超參數(shù)（如后端使用的設(shè)備、每核心批處理大小、預(yù)測長度等）以及從Hugging Face獲取預(yù)訓(xùn)練模型的倉庫id
# 實例化TimesFm模型對象，后續(xù)將利用該模型進行數(shù)據(jù)評估和微調(diào)等操作
## Loading TimesFM pretrained checkpoint
tfm = timesfm.TimesFm(hparams=timesfm.TimesFmHparams(backend="gpu",per_core_batch_size=32,horizon_len=128,),checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id="google/timesfm-1.0-200m"),)# 配置車輛銷售數(shù)據(jù)集相關(guān)信息，包括數(shù)據(jù)集劃分邊界、數(shù)據(jù)文件路徑、時間頻率等，用于后續(xù)的數(shù)據(jù)加載和處理
# 此處簡化假設(shè)按年劃分?jǐn)?shù)據(jù)集，前3年作為訓(xùn)練集，第4年作為驗證集，第5年作為測試集，時間頻率設(shè)置為按年（"Y"）
# 根據(jù)實際數(shù)據(jù)情況和需求，這些設(shè)置都可以進行相應(yīng)調(diào)整
## Evaluating pretrained checkpoint on vehicle sales dataset
DATA_DICT = {"vehicle_sales": {"boundaries": [3, 4, 5],  # 簡化按年劃分，前3年訓(xùn)練，第4年驗證，第5年測試"data_path": "../datasets/vehicle_sales.csv","freq": "Y",  # 按年的時間頻率}
}
dataset = "vehicle_sales"
data_path = DATA_DICT[dataset]["data_path"]
freq = DATA_DICT[dataset]["freq"]
int_freq = timesfm.freq_map(freq)
boundaries = DATA_DICT[dataset]["boundaries"]# 讀取車輛銷售數(shù)據(jù)文件為DataFrame格式，后續(xù)將基于此數(shù)據(jù)進行進一步處理，例如劃分?jǐn)?shù)據(jù)集、提取特征列等操作
# 需要確保數(shù)據(jù)文件路徑正確以及數(shù)據(jù)格式能被read_csv方法正常讀取，如有特殊格式可按需調(diào)整參數(shù)（如分隔符、編碼等）
data_df = pd.read_csv(open(data_path, "r"))# 定義時間序列特征列，這里選取銷售量和銷售價格作為時間序列特征，將用于模型的輸入和預(yù)測等相關(guān)操作
# 根據(jù)實際業(yè)務(wù)需求和數(shù)據(jù)特點，可調(diào)整此列表包含的列名
ts_cols = ["sales_volume", "sales_price"]  
# 暫未定義數(shù)值型協(xié)變量列，可根據(jù)后續(xù)是否需要添加額外數(shù)值型特征進行設(shè)置
num_cov_cols = None  
# 定義分類特征列，這里選取車型和經(jīng)銷商作為分類特征，模型可以根據(jù)這些特征學(xué)習(xí)不同分類下的銷售模式等信息
cat_cov_cols = ["car_model", "dealer"]  context_len = 512
pred_len = 96num_ts = len(ts_cols)
batch_size = 16# 實例化TimeSeriesdata類，用于加載、劃分和預(yù)處理車輛銷售數(shù)據(jù)，配置訓(xùn)練集、驗證集、測試集的范圍，以及設(shè)置數(shù)據(jù)歸一化等參數(shù)
# 該類內(nèi)部會根據(jù)設(shè)置對數(shù)據(jù)進行相應(yīng)處理，生成對應(yīng)的批次數(shù)據(jù)，便于后續(xù)模型訓(xùn)練和評估使用
dtl = data_loader.TimeSeriesdata(data_path=data_path,datetime_col="date",num_cov_cols=num_cov_cols,cat_cov_cols=cat_cov_cols,ts_cols=np.array(ts_cols),train_range=[0, boundaries[0]],val_range=[boundaries[0], boundaries[1]],test_range=[boundaries[1], boundaries[2]],hist_len=context_len,pred_len=pred_len,batch_size=num_ts,freq=freq,normalize=True,epoch_len=None,holiday=False,permute=True,)
# 獲取訓(xùn)練集批次數(shù)據(jù)，每個批次的數(shù)據(jù)將按照設(shè)置的batch_size進行劃分，便于在訓(xùn)練循環(huán)中迭代使用
train_batches = dtl.tf_dataset(mode="train", shift=1).batch(batch_size)
# 獲取驗證集批次數(shù)據(jù)，同樣按照設(shè)置進行劃分，用于在模型訓(xùn)練過程中的定期驗證，以監(jiān)控模型性能和防止過擬合等
val_batches = dtl.tf_dataset(mode="val", shift=pred_len)
# 獲取測試集批次數(shù)據(jù)，用于最終評估模型在未見過的數(shù)據(jù)上的性能表現(xiàn)
test_batches = dtl.tf_dataset(mode="test", shift=pred_len)
# 簡單遍歷訓(xùn)練集批次數(shù)據(jù)的迭代器，此處主要是為了觸發(fā)數(shù)據(jù)加載等相關(guān)操作，確保數(shù)據(jù)可以正常獲取，暫未對數(shù)據(jù)做具體處理
for tbatch in tqdm(train_batches.as_numpy_iterator()):pass
# 打印訓(xùn)練集批次數(shù)據(jù)中第一個元素（通常是輸入數(shù)據(jù)部分）的形狀，用于檢查數(shù)據(jù)維度是否符合預(yù)期
print(tbatch[0].shape)# 以下代碼塊用于計算預(yù)訓(xùn)練模型在測試集上的平均絕對誤差（MAE），通過迭代測試集批次數(shù)據(jù)
# 利用預(yù)訓(xùn)練模型進行預(yù)測，然后對比預(yù)測值和實際值計算平均絕對誤差，以此評估模型初始性能
### MAE on the test split for the pretrained TimesFM model
mae_losses = []
for batch in tqdm(test_batches.as_numpy_iterator()):past = batch[0]actuals = batch[3]forecasts, _ = tfm.forecast(list(past), [0] * past.shape[0], normalize=True)forecasts = forecasts[:, 0 : actuals.shape[1]]mae_losses.append(np.abs(forecasts - actuals).mean())print(f"MAE: {np.mean(mae_losses)}")# 導(dǎo)入一系列用于構(gòu)建和訓(xùn)練模型的praxis、paxml相關(guān)的模塊和類，這些模塊提供了配置模型、定義學(xué)習(xí)器、優(yōu)化訓(xùn)練過程等功能
# 后續(xù)將利用這些工具來對模型進行微調(diào)操作，使其能更好地適應(yīng)車輛銷售數(shù)據(jù)特點和預(yù)測任務(wù)
## Finetuning the model on the vehicle sales dataset
import jax
from jax import numpy as jnp
from praxis import pax_fiddle
from praxis import py_utils
from praxis import pytypes
from praxis import base_model
from praxis import optimizers
from praxis import schedules
from praxis import base_hyperparams
from praxis import base_layer
from paxml import tasks_lib
from paxml import trainer_lib
from paxml import checkpoints
from paxml import learners
from paxml import partitioning
from paxml import checkpoint_types
# PAX shortcuts，定義一些便捷使用的類型和函數(shù)別名，方便后續(xù)代碼中調(diào)用相關(guān)功能時書寫簡潔
NestedMap = py_utils.NestedMap
WeightInit = base_layer.WeightInit
WeightHParams = base_layer.WeightHParams
InstantiableParams = py_utils.InstantiableParams
JTensor = pytypes.JTensor
NpTensor = pytypes.NpTensor
WeightedScalars = pytypes.WeightedScalars
instantiate = base_hyperparams.instantiate
LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]
AuxLossStruct = base_layer.AuxLossStructAUX_LOSS = base_layer.AUX_LOSS
template_field = base_layer.template_field# 定義標(biāo)準(zhǔn)的偽隨機數(shù)生成器（PRNG）鍵名稱，用于在模型訓(xùn)練等過程中需要隨機數(shù)的地方保持一致性和可復(fù)現(xiàn)性
# PARAMS和RANDOM是常見的用于區(qū)分不同用途隨機數(shù)的標(biāo)識
PARAMS = base_layer.PARAMS
RANDOM = base_layer.RANDOM# 生成一個初始的隨機數(shù)生成器的鍵，設(shè)置種子為1234，以便后續(xù)在需要隨機初始化等操作時能復(fù)現(xiàn)結(jié)果
key = jax.random.PRNGKey(seed=1234)
# 配置微調(diào)模型的結(jié)構(gòu)，使用PatchedDecoderFinetuneModel作為基礎(chǔ)結(jié)構(gòu)，并將之前加載的預(yù)訓(xùn)練模型的核心層配置傳入
# 以此構(gòu)建微調(diào)模型的初始結(jié)構(gòu)，后續(xù)將在此基礎(chǔ)上進行參數(shù)更新等微調(diào)操作
model = pax_fiddle.Config(patched_decoder.PatchedDecoderFinetuneModel,name='patched_decoder_finetune',core_layer_tpl=tfm.model_p,
)# 定義構(gòu)建學(xué)習(xí)器的函數(shù)，配置學(xué)習(xí)器相關(guān)參數(shù)，如損失函數(shù)名稱、優(yōu)化器（這里使用Adam優(yōu)化器）及其參數(shù)（學(xué)習(xí)率、學(xué)習(xí)率調(diào)度策略、梯度裁剪閾值、指數(shù)移動平均衰減等）
# 同時設(shè)置在微調(diào)過程中要固定的模型層（通過bprop_variable_exclusion參數(shù)指定，這里示例中固定了變壓器層，可根據(jù)實際情況調(diào)整）
### We will hold the transformer layers fixed while finetuning, while training all other components.
@pax_fiddle.auto_config
def build_learner() -> learners.Learner():return pax_fiddle.Config(learners.Learner,name='learner',loss_name='avg_qloss',optimizer=optimizers.Adam(epsilon=1e-7,clip_threshold=1e2,learning_rate=1e-3,  # 示例中適當(dāng)調(diào)整學(xué)習(xí)率，可根據(jù)實際情況進一步優(yōu)化lr_schedule=pax_fiddle.Config(schedules.Cosine,initial_value=1e-4,final_value=1e-5,total_steps=40000,),ema_decay=0.9999,),# 線性探測，固定變壓器層（可根據(jù)實際情況調(diào)整要固定的層）bprop_variable_exclusion=['.*/stacked_transformer_layer/.*'],)# 構(gòu)建訓(xùn)練任務(wù)配置，將之前定義的模型和學(xué)習(xí)器配置組合起來，同時設(shè)置模型相關(guān)的分布式訓(xùn)練的一些參數(shù)（如mesh形狀和軸名稱等）
# 用于后續(xù)在多設(shè)備等分布式環(huán)境下進行模型訓(xùn)練時的配置管理
task_p = tasks_lib.SingleTask(name='vehicle_sales_learn',model=model,train=tasks_lib.SingleTask.Train(learner=build_learner(),),
)
task_p.model.ici_mesh_shape = [1, 1, 1]
task_p.model.mesh_axis_names = ['replica', 'data', 'mdl']# 獲取可用的設(shè)備（如GPU或CPU等）信息，并將其整理為特定的形狀，用于構(gòu)建分布式訓(xùn)練的Mesh對象，以支持模型在多設(shè)備上并行訓(xùn)練
DEVICES = np.array(jax.devices()).reshape([1, 1, 1])
MESH = jax.sharding.Mesh(DEVICES, ['replica', 'data', 'mdl'])# 獲取本地設(shè)備數(shù)量，用于后續(xù)在數(shù)據(jù)劃分、并行操作等方面根據(jù)設(shè)備數(shù)量進行相應(yīng)處理，并打印設(shè)備相關(guān)信息方便查看
num_devices = jax.local_device_count()
print(f'num_devices: {num_devices}')
print(f'device kind: {jax.local_devices()[0].device_kind}')
jax_task = task_p
key, init_key = jax.random.split(key)# 以下兩個函數(shù)用于處理訓(xùn)練批次數(shù)據(jù)和評估批次數(shù)據(jù)，主要功能是對數(shù)據(jù)的形狀進行調(diào)整和整理
# 確保數(shù)據(jù)格式符合模型輸入和后續(xù)處理的要求，例如將數(shù)據(jù)按照設(shè)備數(shù)量和批次大小等進行合理重塑
# 方便在分布式訓(xùn)練和評估過程中正確使用數(shù)據(jù)
# To correctly prepare a batch of data for model initialization (now that shape
# inference is merged), we take one devices*batch_size tensor tuple of data,
# slice out just one batch, then run the prepare_input_batch function over it.
def process_train_batch(batch):past_ts = batch[0].reshape(batch_size * num_ts, -1)actual_ts = batch[3].reshape(batch_size * num_ts, -1)return NestedMap(input_ts=past_ts, actual_ts=actual_ts)def process_eval_batch(batch):past_ts = batch[0]actual_ts = batch[3]return NestedMap(input_ts=past_ts, actual_ts=actual_ts)# 初始化模型狀態(tài)，傳入訓(xùn)練任務(wù)配置、初始化隨機數(shù)鍵以及處理后的訓(xùn)練批次數(shù)據(jù)等信息
# 根據(jù)指定的檢查點類型（這里是GDA類型）進行模型狀態(tài)的初始化操作，得到初始的模型狀態(tài)信息
jax_model_states, _ = trainer_lib.initialize_model_state(jax_task,init_key,process_train_batch(tbatch),checkpoint_type=checkpoint_types.CheckpointType.GDA,
)# 將預(yù)訓(xùn)練模型的參數(shù)設(shè)置為微調(diào)模型的初始權(quán)重，具體是將預(yù)訓(xùn)練模型的參數(shù)賦值給微調(diào)模型狀態(tài)中對應(yīng)核心層的參數(shù)部分
# 這樣微調(diào)模型就可以在預(yù)訓(xùn)練的基礎(chǔ)上進行進一步優(yōu)化，加快收斂并利用預(yù)訓(xùn)練學(xué)到的通用特征表示
### Setting the initial model weights to the pretrained TimesFM parameters.
jax_model_states.mdl_vars['params']['core_layer'] = tfm._train_state.mdl_vars['params']
jax_vars = jax_model_states.mdl_vars
gc.collect()# 以下是模型微調(diào)的訓(xùn)練循環(huán)部分，定義了訓(xùn)練步和評估步的函數(shù)，在多個訓(xùn)練輪次（epoch）內(nèi)循環(huán)進行訓(xùn)練和定期評估
# 通過早停機制（根據(jù)驗證集損失決定是否提前停止訓(xùn)練）來避免過擬合，在每個訓(xùn)練步中更新模型參數(shù)，每個評估步計算驗證集上的損失并保存最優(yōu)模型狀態(tài)的檢查點
### Training loop# 將之前配置好的訓(xùn)練任務(wù)（task_p）賦值給jax_task變量，后續(xù)在訓(xùn)練和評估步驟中會用到這個任務(wù)配置信息
jax_task = task_p# 定義訓(xùn)練步函數(shù)train_step，該函數(shù)內(nèi)部調(diào)用了trainer_lib.train_step_single_learner函數(shù)，用于執(zhí)行單個學(xué)習(xí)器的一次訓(xùn)練步驟
# 它接收當(dāng)前的模型狀態(tài)（states）、偽隨機數(shù)生成器的鍵（prng_key）以及輸入數(shù)據(jù)（inputs）作為參數(shù)，返回訓(xùn)練后的模型狀態(tài)等相關(guān)信息
def train_step(states, prng_key, inputs):return trainer_lib.train_step_single_learner(jax_task, states, prng_key, inputs)# 定義評估步函數(shù)eval_step，首先將模型狀態(tài)轉(zhuǎn)換為評估狀態(tài)（通過to_eval_state方法），然后調(diào)用trainer_lib.eval_step_single_learner函數(shù)執(zhí)行評估步驟
# 同樣接收模型狀態(tài)、偽隨機數(shù)生成器的鍵以及輸入數(shù)據(jù)作為參數(shù)，返回評估相關(guān)的結(jié)果信息（例如損失值等）
def eval_step(states, prng_key, inputs):states = states.to_eval_state()return trainer_lib.eval_step_single_learner(jax_task, states, prng_key, inputs)# 對初始的隨機數(shù)生成器的鍵（key）進行分割，生成三個新的隨機數(shù)生成器的鍵，分別用于后續(xù)的訓(xùn)練、評估以及其他可能的操作
# 這樣可以保證在不同的步驟中使用不同的隨機數(shù)流，便于控制隨機性和復(fù)現(xiàn)實驗結(jié)果
key, train_key, eval_key = jax.random.split(key, 3)# 根據(jù)本地設(shè)備數(shù)量（jax.local_device_count()），將訓(xùn)練用的隨機數(shù)生成器的鍵（train_key）分割成多個子鍵，每個子鍵對應(yīng)一個設(shè)備
# 用于在分布式訓(xùn)練中為每個設(shè)備提供獨立的隨機數(shù)種子，確保隨機性在不同設(shè)備上的正確應(yīng)用
train_prng_seed = jax.random.split(train_key, num=jax.local_device_count())# 同理，將評估用的隨機數(shù)生成器的鍵（eval_key）也按照本地設(shè)備數(shù)量進行分割，為每個設(shè)備在評估過程中提供獨立的隨機數(shù)種子
eval_prng_seed = jax.random.split(eval_key, num=jax.local_device_count())# 使用jax.pmap對訓(xùn)練步函數(shù)（train_step）進行并行化處理，指定按照'batch'軸進行并行，使得訓(xùn)練可以在多個設(shè)備上并行執(zhí)行，提高訓(xùn)練效率
p_train_step = jax.pmap(train_step, axis_name='batch')# 同樣地，對評估步函數(shù)（eval_step）進行并行化處理，使其能在多個設(shè)備上并行執(zhí)行評估操作，也是按照'batch'軸進行并行
p_eval_step = jax.pmap(eval_step, axis_name='batch')# 對初始的模型狀態(tài)（jax_model_states）進行復(fù)制操作，以適配分布式訓(xùn)練環(huán)境，使得每個設(shè)備都有一份相同的初始模型狀態(tài)副本
# 后續(xù)每個設(shè)備可以基于這份副本進行獨立的參數(shù)更新等操作，最終再進行匯總等處理
replicated_jax_states = trainer_lib.replicate_model_state(jax_model_states)# 獲取復(fù)制后的模型狀態(tài)中的變量部分（mdl_vars），方便后續(xù)在訓(xùn)練和評估過程中對模型參數(shù)等變量進行操作和訪問
replicated_jax_vars = replicated_jax_states.mdl_vars# 初始化最優(yōu)驗證集損失值為一個較大的數(shù)（1e7），在訓(xùn)練過程中，一旦發(fā)現(xiàn)更小的驗證集損失值，就會更新這個最優(yōu)值，并保存對應(yīng)的模型狀態(tài)
best_eval_loss = 1e7# 記錄當(dāng)前已經(jīng)執(zhí)行的訓(xùn)練步數(shù)，初始化為0，隨著訓(xùn)練循環(huán)的進行，每執(zhí)行一次訓(xùn)練步就會加1，用于判斷是否達(dá)到定期評估的步數(shù)等條件
step_count = 0# 早停機制相關(guān)的耐心值（patience），初始化為0，代表目前還沒有出現(xiàn)驗證集損失不再下降的情況
# 當(dāng)驗證集損失連續(xù)多次（由PATIENCE變量定義）沒有下降時，就會觸發(fā)早停機制，提前結(jié)束訓(xùn)練
patience = 0# 設(shè)定總的訓(xùn)練輪次（epoch）數(shù)量，這里設(shè)置為100，表示模型將對整個訓(xùn)練數(shù)據(jù)集完整遍歷100次，可根據(jù)實際情況調(diào)整該值
NUM_EPOCHS = 100# 設(shè)定早停機制中的耐心值，即驗證集損失連續(xù)多少次沒有下降就觸發(fā)早停，這里設(shè)置為5，意味著如果連續(xù)5次評估驗證集損失都沒有變小，就停止訓(xùn)練
PATIENCE = 5# 定義每經(jīng)過多少個訓(xùn)練步就進行一次模型在驗證集上的評估操作，這里設(shè)置為1000步評估一次，用于定期監(jiān)控模型在驗證集上的性能表現(xiàn)
TRAIN_STEPS_PER_EVAL = 1000# 指定保存模型檢查點的目錄路徑，訓(xùn)練過程中，當(dāng)發(fā)現(xiàn)當(dāng)前模型在驗證集上的性能更好（驗證集損失更小）時，會將模型狀態(tài)保存到這個目錄下
CHECKPOINT_DIR = '/home/senrajat_google_com/vehicle_sales_finetune'# 定義一個函數(shù)reshape_batch_for_pmap，用于根據(jù)設(shè)備數(shù)量對批次數(shù)據(jù)進行形狀重塑，使其能正確地在分布式訓(xùn)練環(huán)境下分配到各個設(shè)備上
# 具體操作是將輸入張量的第一個維度（通常是批次大小維度）按照設(shè)備數(shù)量進行劃分，重新調(diào)整張量的形狀
def reshape_batch_for_pmap(batch, num_devices):def _reshape(input_tensor):bsize = input_tensor.shape[0]residual_shape = list(input_tensor.shape[1:])nbsize = bsize // num_devicesreturn jnp.reshape(input_tensor, [num_devices, nbsize] + residual_shape)return jax.tree.map(_reshape, batch)# 外層循環(huán)，按照設(shè)定的訓(xùn)練輪次（NUM_EPOCHS）進行循環(huán)訓(xùn)練，每個epoch代表對整個訓(xùn)練數(shù)據(jù)集的一次完整遍歷
for epoch in range(NUM_EPOCHS):# 打印當(dāng)前所處的訓(xùn)練輪次信息，方便在訓(xùn)練過程中查看訓(xùn)練進度，flush=True用于立即輸出信息，不緩沖print(f"__________________Epoch: {epoch}__________________", flush=True)# 獲取訓(xùn)練集批次數(shù)據(jù)的迭代器，用于在當(dāng)前epoch內(nèi)逐個批次地遍歷訓(xùn)練數(shù)據(jù)train_its = train_batches.as_numpy_iterator()# 判斷如果早停的耐心值（patience）達(dá)到設(shè)定的閾值（PATIENCE），則觸發(fā)早停機制，結(jié)束訓(xùn)練if patience >= PATIENCE:print("Early stopping.", flush=True)break# 內(nèi)層循環(huán)，遍歷當(dāng)前epoch的每個訓(xùn)練批次數(shù)據(jù)for batch in tqdm(train_its):train_losses = []# 再次檢查早停條件，若滿足則提前停止當(dāng)前批次的訓(xùn)練if patience >= PATIENCE:print("Early stopping.", flush=True)break# 調(diào)用函數(shù)處理訓(xùn)練批次數(shù)據(jù)，主要是對數(shù)據(jù)形狀進行調(diào)整，使其符合模型訓(xùn)練輸入要求tbatch = process_train_batch(batch)# 根據(jù)設(shè)備數(shù)量對處理后的批次數(shù)據(jù)進行重塑，以便在分布式訓(xùn)練環(huán)境下能正確分配到各個設(shè)備上進行并行計算tbatch = reshape_batch_for_pmap(tbatch, num_devices)# 執(zhí)行分布式訓(xùn)練的一個訓(xùn)練步，傳入當(dāng)前模型狀態(tài)、訓(xùn)練隨機數(shù)種子以及處理好的批次數(shù)據(jù)# 返回更新后的模型狀態(tài)以及包含訓(xùn)練損失等信息的輸出結(jié)果（step_fun_out）replicated_jax_states, step_fun_out = p_train_step(replicated_jax_states, train_prng_seed, tbatch)# 將當(dāng)前訓(xùn)練步的損失值添加到訓(xùn)練損失列表（train_losses）中，后續(xù)可以用于計算平均訓(xùn)練損失等操作train_losses.append(step_fun_out.loss[0])# 判斷當(dāng)前訓(xùn)練步數(shù)是否達(dá)到了設(shè)定的定期評估步數(shù)（TRAIN_STEPS_PER_EVAL），如果達(dá)到則進行模型在驗證集上的評估操作if step_count % TRAIN_STEPS_PER_EVAL == 0:# 打印當(dāng)前訓(xùn)練步數(shù)下的平均訓(xùn)練損失值，方便查看訓(xùn)練過程中模型在訓(xùn)練集上的損失變化情況，flush=True用于立即輸出信息print(f"Train loss at step {step_count}: {np.mean(train_losses)}",flush=True,)# 清空訓(xùn)練損失列表，為下一個評估周期準(zhǔn)備，避免累計之前的損失值影響下一次平均損失的計算train_losses = []# 打印提示信息，表示開始進行模型在驗證集上的評估操作print("Starting eval.", flush=True)# 獲取驗證集批次數(shù)據(jù)的迭代器，用于在驗證過程中逐個批次地遍歷驗證數(shù)據(jù)val_its = val_batches.as_numpy_iterator()# 初始化用于存儲每個驗證批次損失值的列表，用于后續(xù)計算平均驗證集損失eval_losses = []# 遍歷驗證集的每個批次數(shù)據(jù)for ev_batch in tqdm(val_its):# 調(diào)用函數(shù)處理驗證批次數(shù)據(jù)，對數(shù)據(jù)形狀進行調(diào)整，使其符合模型評估輸入要求ebatch = process_eval_batch(ev_batch)# 根據(jù)設(shè)備數(shù)量對處理后的驗證批次數(shù)據(jù)進行重塑，適配分布式評估環(huán)境ebatch = reshape_batch_for_pmap(ebatch, num_devices)# 執(zhí)行分布式評估的一個評估步，傳入當(dāng)前模型狀態(tài)、評估隨機數(shù)種子以及處理好的驗證批次數(shù)據(jù)# 返回包含評估損失等信息的輸出結(jié)果（這里只關(guān)心損失值，所以用下劃線忽略其他返回信息）_, step_fun_out = p_eval_step(replicated_jax_states, eval_prng_seed, ebatch)# 將當(dāng)前驗證批次的損失值添加到驗證損失列表（eval_losses）中eval_losses.append(step_fun_out.loss[0])# 計算平均驗證集損失值，通過對驗證損失列表中的所有損失值求平均得到mean_loss = np.mean(eval_losses)# 打印當(dāng)前訓(xùn)練步數(shù)下的平均驗證集損失值，方便查看模型在驗證集上的性能表現(xiàn)，flush=True用于立即輸出信息print(f"Eval loss at step {step_count}: {mean_loss}", flush=True)# 判斷當(dāng)前平均驗證集損失值是否小于之前記錄的最優(yōu)驗證集損失值（best_eval_loss），或者是否為NaN（表示出現(xiàn)異常情況）# 如果滿足條件，說明當(dāng)前模型在驗證集上的性能更好，需要保存當(dāng)前的模型狀態(tài)作為最優(yōu)模型狀態(tài)if mean_loss < best_eval_loss or np.isnan(mean_loss):# 更新最優(yōu)驗證集損失值為當(dāng)前的平均驗證集損失值best_eval_loss = mean_loss# 打印提示信息，表示正在保存模型檢查點print("Saving checkpoint.")# 對復(fù)制后的模型狀態(tài)進行處理，將其轉(zhuǎn)換為適合保存的格式（可能涉及去除一些分布式相關(guān)的冗余信息等操作）jax_state_for_saving = py_utils.maybe_unreplicate_for_fully_replicated(replicated_jax_states)# 調(diào)用函數(shù)保存模型檢查點，將處理后的模型狀態(tài)保存到指定的目錄（CHECKPOINT_DIR）下，并且設(shè)置覆蓋已存在的同名檢查點checkpoints.save_checkpoint(jax_state_for_saving, CHECKPOINT_DIR, overwrite=True)# 將早停機制的耐心值重置為0，因為當(dāng)前模型性能有提升，重新開始計算耐心值patience = 0# 刪除已經(jīng)保存的模型狀態(tài)變量，釋放內(nèi)存空間，避免內(nèi)存占用過多del jax_state_for_saving# 手動觸發(fā)垃圾回收，及時回收不再使用的內(nèi)存，優(yōu)化內(nèi)存使用情況gc.collect()# 如果當(dāng)前平均驗證集損失值沒有小于最優(yōu)驗證集損失值，說明模型在驗證集上的性能沒有提升，則增加早停機制的耐心值else:patience += 1# 打印當(dāng)前的耐心值，方便查看早停機制的觸發(fā)進度情況print(f"patience: {patience}")# 每執(zhí)行完一個訓(xùn)練步，訓(xùn)練步數(shù)加1，用于跟蹤訓(xùn)練的進度以及判斷是否達(dá)到定期評估等條件step_count += 1# 以下代碼用于加載根據(jù)驗證集損失選出的最優(yōu)微調(diào)后的模型檢查點，并在測試集上對其進行評估，計算平均絕對誤差（MAE）來衡量模型性能
## Loading and evaluating the best (according to validation loss) finetuned checkpoint# 調(diào)用函數(shù)從指定的目錄（CHECKPOINT_DIR）中恢復(fù)之前保存的最優(yōu)模型狀態(tài)，將其賦值給train_state變量，用于后續(xù)的模型參數(shù)更新和評估操作
train_state = checkpoints.restore_checkpoint(jax_model_states, CHECKPOINT_DIR)# 打印恢復(fù)的模型狀態(tài)對應(yīng)的訓(xùn)練步數(shù)信息，可用于查看加載的是哪個階段保存的模型狀態(tài)
print(train_state.step)# 將微調(diào)后模型的參數(shù)更新到原TimesFM模型中，具體是將恢復(fù)的模型狀態(tài)中的核心層參數(shù)賦值給原TimesFM模型的對應(yīng)參數(shù)部分
# 使得原模型可以使用微調(diào)后的參數(shù)進行預(yù)測等操作，用于在測試集上評估微調(diào)后的模型性能
tfm._train_state.mdl_vars['params'] = train_state.mdl_vars['params']['core_layer']# 對TimesFM模型執(zhí)行即時編譯（jit）相關(guān)的解碼操作，可能是為了優(yōu)化模型在后續(xù)預(yù)測過程中的性能，加快預(yù)測速度
tfm.jit_decode()# 初始化用于存儲測試集上每個批次預(yù)測結(jié)果與實際結(jié)果的平均絕對誤差（MAE）的列表，后續(xù)將通過循環(huán)計算并填充該列表
mae_losses = []# 遍歷測試集的每個批次數(shù)據(jù)，用于計算在整個測試集上的平均絕對誤差（MAE）
for batch in tqdm(test_batches.as_numpy_iterator()):# 獲取當(dāng)前批次的輸入數(shù)據(jù)（通常是歷史時間序列數(shù)據(jù)等），作為模型預(yù)測的輸入past = batch[0]# 獲取當(dāng)前批次的實際值（真實的目標(biāo)數(shù)據(jù)，例如實際銷售量、銷售價格等），用于與模型預(yù)測結(jié)果進行對比計算誤差actuals = batch[3]# 使用更新參數(shù)后的TimesFM模型進行預(yù)測，傳入當(dāng)前批次的輸入數(shù)據(jù)以及一些相關(guān)的輔助參數(shù)（這里輔助參數(shù)都設(shè)置為0，具體含義可能取決于模型的定義）# 返回預(yù)測結(jié)果（forecasts）以及其他可能的相關(guān)信息（這里用下劃線忽略）_, forecasts = tfm.forecast(list(past), [0] * past.shape[0])# 對預(yù)測結(jié)果進行維度處理，選取與實際值維度對應(yīng)的部分，確保兩者可以正確地進行誤差計算（這里假設(shè)實際值和預(yù)測值在維度上需要進行一定的對齊操作）forecasts = forecasts[:, 0 : actuals.shape[1], 5]# 計算當(dāng)前批次預(yù)測結(jié)果與實際結(jié)果的平均絕對誤差（MAE），通過計算預(yù)測值與實際值差值的絕對值的平均值得到# 將每個批次的MAE值添加到mae_losses列表中，后續(xù)可以通過求平均得到整個測試集上的平均MAE值mae_losses.append(np.abs(forecasts - actuals).mean())print(f"MAE: {np.mean(mae_losses)}")

請注意：

1. 上述代碼中的路徑等相關(guān)設(shè)置（如?CHECKPOINT_DIR、數(shù)據(jù)文件路徑等）需要根據(jù)你的實際運行環(huán)境進行調(diào)整，確保可以正確讀寫文件以及保存和加載模型檢查點。
2. 代碼中關(guān)于模型的一些超參數(shù)（如學(xué)習(xí)率、訓(xùn)練輪數(shù)、批處理大小等）都是示例值，你可能需要根據(jù)車輛銷售數(shù)據(jù)的特點、模型訓(xùn)練情況等進行多次試驗和調(diào)整，以獲得更好的預(yù)測性能。
3. 假設(shè)數(shù)據(jù)文件?vehicle_sales.csv?的格式是比較規(guī)范的，能被?pandas?的?read_csv?方法正常讀取，如果實際數(shù)據(jù)有特殊格式（例如包含標(biāo)題行、特定的日期格式、缺失值表示等情況），可能需要進一步對數(shù)據(jù)讀取部分進行修改完善。