Python 按列（对象）分层拆分

当尝试按列（分类）进行Strafted拆分时，它返回错误

Country     ColumnA    ColumnB   ColumnC   Label
AB            0.2        0.5       0.1       14  
CD            0.9        0.2       0.6       60
EF            0.4        0.3       0.8       5
FG            0.6        0.9       0.2       15  

这是我的密码：

X = df.loc[:, df.columns != 'Label']
y = df['Label']

# Train/test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0, stratify=df.Country)

from sklearn.linear_model import LinearRegression
lm = LinearRegression()
lm.fit(X_train,y_train)
lm_predictions = lm.predict(X_test)

因此，我得到如下错误：

ValueError: could not convert string to float: 'AB'

---

在复制代码时，我发现错误来自于试图在一组包含字符串的特性上拟合线性回归模型。为您提供了一些操作选项。我建议使用

X\u-train，X\u-test=pd.get\u-dummies（X\u-train.Country），pd.get\u-dummies（X\u-test.Country）

---

在你进行train_test_split（）之后对你的国家进行一次热编码，以保持你正在寻找的类平衡。

在复制代码时，我发现错误来自于试图在一组包含字符串的特征上拟合线性回归模型。为您提供了一些操作选项。我建议使用

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

df = pd.DataFrame({
        'Country': ['AB', 'CD', 'EF', 'FG']*20,
        'ColumnA' : [1]*20*4,'ColumnB' : [10]*20*4, 'Label': [1,0,1,0]*20
    })

df['Country_Code'] = df['Country'].astype('category').cat.codes

X = df.loc[:, df.columns.drop(['Label','Country'])]
y = df['Label']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0, stratify=df.Country_Code)
lm = LinearRegression()
lm.fit(X_train,y_train)
lm_predictions = lm.predict(X_test)

X\u-train，X\u-test=pd.get\u-dummies（X\u-train.Country），pd.get\u-dummies（X\u-test.Country）

在你进行训练测试后，对你的国家进行一次热编码，以保持你所寻找的阶级平衡

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

df = pd.DataFrame({
        'Country': ['AB', 'CD', 'EF', 'FG']*20,
        'ColumnA' : [1]*20*4,'ColumnB' : [10]*20*4, 'Label': [1,0,1,0]*20
    })

df['Country_Code'] = df['Country'].astype('category').cat.codes

X = df.loc[:, df.columns.drop(['Label','Country'])]
y = df['Label']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0, stratify=df.Country_Code)
lm = LinearRegression()
lm.fit(X_train,y_train)
lm_predictions = lm.predict(X_test)

• 将

  country

  中的字符串值转换为数字，并将其另存为新列
• 创建

  x

  列车数据下拉列表

  标签
  （
  y
  ）以及字符串
  country
  列时

方法2 如果您将要进行预测的测试数据稍后会出现，则您需要一种机制，在进行预测之前将其

国家

转换为

代码

。在这种情况下，建议使用

LabelEncoder

，您可以使用

fit

方法将字符串编码为标签，然后使用

transform

对测试数据所在国进行编码

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn import preprocessing

df = pd.DataFrame({
        'Country': ['AB', 'CD', 'EF', 'FG']*20,
        'ColumnA' : [1]*20*4,'ColumnB' : [10]*20*4, 'Label': [1,0,1,0]*20
    })

# Train-Validation 
le = preprocessing.LabelEncoder()
df['Country_Code'] = le.fit_transform(df['Country'])
X = df.loc[:, df.columns.drop(['Label','Country'])]
y = df['Label']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0, stratify=df.Country_Code)
lm = LinearRegression()
lm.fit(X_train,y_train)

# Test
test_df = pd.DataFrame({'Country': ['AB'], 'ColumnA' : [1],'ColumnB' : [10] })
test_df['Country_Code'] = le.transform(test_df['Country'])
print (lm.predict(test_df.loc[:, test_df.columns.drop(['Country'])]))

• 将

  country

  中的字符串值转换为数字，并将其另存为新列
• 创建

  x

  列车数据下拉列表

  标签
  （
  y
  ）以及字符串
  country
  列时

方法2 如果您将要进行预测的测试数据稍后会出现，则您需要一种机制，在进行预测之前将其

国家

转换为

代码

。在这种情况下，建议使用

LabelEncoder

，您可以使用

fit

方法将字符串编码为标签，然后使用

transform

对测试数据所在国进行编码

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn import preprocessing

df = pd.DataFrame({
        'Country': ['AB', 'CD', 'EF', 'FG']*20,
        'ColumnA' : [1]*20*4,'ColumnB' : [10]*20*4, 'Label': [1,0,1,0]*20
    })

# Train-Validation 
le = preprocessing.LabelEncoder()
df['Country_Code'] = le.fit_transform(df['Country'])
X = df.loc[:, df.columns.drop(['Label','Country'])]
y = df['Label']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=0, stratify=df.Country_Code)
lm = LinearRegression()
lm.fit(X_train,y_train)

# Test
test_df = pd.DataFrame({'Country': ['AB'], 'ColumnA' : [1],'ColumnB' : [10] })
test_df['Country_Code'] = le.transform(test_df['Country'])
print (lm.predict(test_df.loc[:, test_df.columns.drop(['Country'])]))

---

无法重现错误（使用“国家/地区”作为“国家/地区代码”）@ChristianSloper正确点，已修复。Thanks@LucaMassaron你能帮忙吗？Thank无法重现错误（使用“国家”作为“国家代码”）@ChristianSloper正确点，已修复。Thanks@LucaMassaron你能帮忙吗？谢谢