请你把这段代码用def的形式改写data = torch.randn(16992, 307, 12, 2) week_feature = torch.zeros((data.shape[0], 1)) # 遍历所有时间点，获取对应的星期并将星期转化为对应的数字 for i in range(data.shape[0]): date_str = '2016-01-01 00:00:00' # 假设数据集中的时间从2016年1月1日开始 date = datetime.datetime.strptime(date_str, '%Y-%m-%d %H:%M:%S') date += datetime.timedelta(minutes=5*i) # 每个时间点间隔5分钟 week = date.weekday() # 获取星期，0表示星期一，1表示星期二，以此类推 week_feature[i] = week + 1 # 将星期转化为对应的数字并将其存入星期特征张量中 # 将星期特征添加到原有的数据集中 week_feature=week_feature.unsqueeze(1).repeat(1, 307, 12, 1) week_feature=week_feature.reshape(16992,307,12,1) data = torch.cat((data, week_feature), dim=-1) # week_feature = week_feature.repeat(1, 307, 12, 1) # 将星期特征张量与数据集在最后一个维度上拼接 # data = torch.cat((data, week_feature), dim=-1) # 输出结果 print(data) print(data.shape) pytorch版本

可以这样改写：

import torch
import datetime

def add_week_feature(data):
    week_feature = torch.zeros((data.shape[0], 1))  
    for i in range(data.shape[0]):
        date_str = '2016-01-01 00:00:00'
        date = datetime.datetime.strptime(date_str, '%Y-%m-%d %H:%M:%S')
        date += datetime.timedelta(minutes=5*i)
        week = date.weekday()
        week_feature[i] = week + 1
    week_feature = week_feature.unsqueeze(1).repeat(1, 307, 12, 1)
    week_feature = week_feature.reshape(16992,307,12,1)
    data = torch.cat((data, week_feature), dim=-1)
    return data

data = torch.randn(16992, 307, 12, 2)
data = add_week_feature(data)
print(data)
print(data.shape)

这样就可以通过调用 `add_week_feature(data)` 函数来实现原有代码的功能。