1. pytorch加载数据并进行训练测试
import time
from torch import optim, nn
from torch.utils import data
from dataset import MyDataset
from model import Net
from function import *
# 加载数据
cur_path = os.getcwd()
data_path = os.path.join(cur_path, "data\preliminary")
train_dataset = MyDataset(mode='train', data_path=data_path)
train_dataloader = data.DataLoader(train_dataset, batch_size=20, shuffle=True)
valid_dataset = MyDataset(mode='valid', data_path=data_path)
valid_dataloader = data.DataLoader(valid_dataset, batch_size=20)
# 实例化模型
model = Net()
# 优化器
optimizer = optim.Adam(model.parameters(), lr=0.0003)
# 损失函数
criterion = nn.CrossEntropyLoss()
# 训练批次
epochs = 50
# 模型保存路径
save_path = 'checkpoint/'
if not os.path.exists(save_path):
os.mkdir(save_path)
# 训练 & 测试过程
train_acc, valid_acc, train_losses, valid_losses = [], [], [], []
best_acc = 0.0
for epoch in range(epochs):
epoch_start = time.time()
model.train()
train_loss = 0.0
train_a = 0.0
valid_loss = 0.0
valid_a = 0.0
t0 = 0
t1 = 0
for i, (inputs, labels) in enumerate(train_dataloader):
inputs, labels = torch.tensor(inputs, dtype=torch.float), torch.tensor(labels, dtype=torch.float)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
train_loss += loss.item()
loss.backward()
optimizer.step()
ret, predictions = torch.max(outputs.data, 1) # 返回概率大的值和索引
labels = labels[:, -1]
acc = torch.sum(predictions == labels) / outputs.shape[0]
train_a += acc.item()
t0 = t0 + 1
with torch.no_grad():
model.eval()
for j, (inputs, labels) in enumerate(valid_dataloader):
inputs, labels = torch.tensor(inputs, dtype=torch.float), torch.tensor(labels, dtype=torch.float)
outputs = model(inputs)
loss = criterion(outputs, labels)
valid_loss += loss.item()
ret, predictions = torch.max(outputs.data, 1)
labels = labels[:, -1]
acc = torch.sum(predictions == labels) / outputs.shape[0]
valid_a += acc.item()
t1 = t1 + 1
train_loss = train_loss / t0
train_a = train_a / t0
valid_loss = valid_loss / t1
valid_a = valid_a / t1
train_acc.append(train_a)
valid_acc.append(valid_a)
train_losses.append(train_loss)
valid_losses.append(valid_loss)
epoch_end = time.time()
# 保存模型
if valid_a > best_acc:
torch.save(model.state_dict(), save_path + 'best_model')
best_acc = valid_a
if epoch == epochs - 1:
torch.save(model.state_dict(), save_path + 'final_model')
print("Epoch: {}/{}, Training:tLoss: {:.4f}, Accuracy: {:.2f}%, "
"ttValidation:tLoss: {:.4f}, Accuracy: {:.2f}%, Time: {:.4f}s".format(
epoch + 1, epochs, train_loss, train_a * 100, valid_loss, valid_a * 100,
epoch_end - epoch_start))
print(f'training end.best_model save to {save_path}.')
plot_acc(train_acc, valid_acc)
plot_loss(train_losses, valid_losses)
plot_results(epochs, train_acc, train_losses, valid_acc, train_losses)
2. 自定义Dataset
from scipy.io import loadmat
import os
from torch.utils import data
import pandas as pd
import numpy as np
def convert2oneHot(index, Lens):
hot = np.zeros((Lens,))
hot[index] = 1
return hot
def normalize(v):
part1 = v - v.mean(axis=1).reshape((v.shape[0], 1))
part2 = v.max(axis=1).reshape((v.shape[0], 1)) + 2e-12
return part1 / part2
class MyDataset(data.Dataset):
def __init__(self, mode, data_path):
super(MyDataset, self).__init__()
self.csv_path = os.path.join(data_path, "reference.csv")
self.data_path = os.path.join(data_path, "TRAIN")
self.temp_list = [] # 文件路径
self._parse_dataset()
self.mode = mode.lower()
if self.mode == 'train':
self.temp_list = self.temp_list[:500]
elif self.mode == 'valid':
self.temp_list = self.temp_list[500:]
else:
raise ValueError('mode must be "train" or "valid"!')
def __getitem__(self, item):
feature = self.get_feature(self.temp_list[item, 0])
label = convert2oneHot(self.temp_list[item, 1], 2)
return feature, label
def __len__(self):
return len(self.temp_list)
def get_feature(self, name):
mat = loadmat(os.path.join(self.data_path, name))
dat = mat['data']
feature = dat[0:12] # feature: (12, 5000)
# return normalize(feature).transpose() # feature: (5000, 12)
return normalize(feature)
def _parse_dataset(self):
self.temp_list = np.array(pd.read_csv(self.csv_path))
3. 创建模型
from torch import nn
# input: (20, 12, 5000)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv1d(in_channels=12, out_channels=16, kernel_size=16, stride=2, padding=8),
nn.ReLU(),
nn.Conv1d(in_channels=16, out_channels=16, kernel_size=16, stride=2, padding=8),
nn.ReLU(),
nn.MaxPool1d(2)
)
self.layer2 = nn.Sequential(
nn.Conv1d(in_channels=16, out_channels=64, kernel_size=8, stride=2, padding=4),
nn.ReLU(),
nn.Conv1d(in_channels=64, out_channels=64, kernel_size=8, stride=2, padding=4),
nn.ReLU(),
nn.MaxPool1d(2)
)
self.layer3 = nn.Sequential(
nn.Conv1d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=2),
nn.ReLU(),
nn.Conv1d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=2),
nn.ReLU(),
nn.MaxPool1d(2)
)
self.layer4 = nn.Sequential(
nn.Conv1d(in_channels=128, out_channels=256, kernel_size=2, stride=1, padding=1),
nn.ReLU(),
nn.Conv1d(in_channels=256, out_channels=256, kernel_size=2, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool1d(2)
)
self.layer5 = nn.Sequential(
nn.AdaptiveAvgPool1d(2),
nn.Flatten()
)
self.layer6 = nn.Sequential(
nn.Linear(in_features=256 * 2, out_features=2),
nn.Dropout(0.3),
nn.Softmax()
)
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.layer5(x)
x = self.layer6(x)
return x
4.加载模型进行预测
# 加载数据
cur_path = os.getcwd()
data_path = os.path.join(cur_path, "data\preliminary\TEST")
files = glob(data_path + '/*.mat')
# 加载模型
model = Net()
model.load_state_dict(torch.load(os.path.join(cur_path, 'checkpoint\best_model')))
model.eval()
predictions = []
file_list = []
for file in files:
inputs = load_data(file)
inputs = torch.tensor(inputs, dtype=torch.float)
inputs = Variable(torch.unsqueeze(inputs, dim=0).float())
pred = model(inputs)
ret, pred_ = torch.max(pred.data, 1)
name = os.path.basename(file)
file_list.append(name)
predictions.append(pred_)
# name = file.split('\')[-1]
print(f'{name} t {pred_.item()}')
save_txt(file_list, predictions, 'pred_result.txt')
