torch如何使用自己的训练集和测试集
时间: 2023-08-10 14:07:12 浏览: 40
使用PyTorch训练自己的数据集需要进行以下步骤:
1. 创建自己的数据集类(继承自`torch.utils.data.Dataset`),实现`__len__`和`__getitem__`方法。
2. 使用`torch.utils.data.DataLoader`加载数据集,并指定batch size、shuffle等参数。
3. 定义模型和损失函数,选择优化器。
4. 在训练循环中,对每个batch的数据进行前向传播、计算损失、反向传播和优化。
5. 在测试循环中,对测试集的数据进行前向传播,计算损失和性能指标。
具体代码如下:
```python
import torch
from torch.utils.data import DataLoader, Dataset
# 1. 创建自己的数据集类
class MyDataset(Dataset):
def __init__(self, data_path):
self.data = torch.load(data_path) # 数据集格式为tensor
self.labels = torch.load(label_path) # 标签格式为tensor
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx], self.labels[idx]
# 2. 加载数据集
train_dataset = MyDataset(train_data_path)
train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_dataset = MyDataset(test_data_path)
test_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=False)
# 3. 定义模型和损失函数、优化器
model = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 4. 训练循环
for epoch in range(num_epochs):
for i, (inputs, labels) in enumerate(train_dataloader):
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 5. 测试循环
test_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_dataloader:
outputs = model(inputs)
loss = criterion(outputs, labels)
test_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
test_accuracy = 100 * correct / total
```