深度学习pytorch整体流程

一、数据处理

复制代码

from torch.utils.data import Dataset, DataLoader
import torch
import numpy as np
import pandas as pd
import os

class custom_dataset(Dataset):
    def __init__(self, root_path, file_path, transform, target_transform):
        self.data = pd.read_csv(os.path.join(root_path, file_path))
        self.transform = transform
        self.target_transform = target_transform
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        row = self.data.iloc[idx]
        data = row['data']
        target = row['target']
        if self.transform:
            data = self.transform(data)
        if self.target_transform:
            target = self.target_transform(target)
        return data, target

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

二、模型建立、训练、测试

复制代码

import torch
from torch import nn

class network(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.linear = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim),
            nn.ReLU()
        )
    def forward(self, x):
        x = self.linear(x)
        return x

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

model = network(input_dim=8, hidden_dim=16, output_dim=8).to(torch.device('cpu'))

learning_rate = 1e-3
batch_size = 64
epochs = 5

loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

def train(dataloader, model, loss_fn, optimizer, device):
    model.train()
    for batch, (x, y) in enumerate(dataloader):
        x, y = x.to(device), y.to(device)
        pred = model(x)
        loss = loss_fn(pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if batch % 100 == 0:
            loss = loss.item()
            print(f'train_loss:{loss}')


def test(dataloader, model, loss_fn):
    num_batches = len(dataloader)
    model.eval()
    with torch.no_grad():
        for x, y in dataloader:
            pred = model(x)
            test_loss += loss_fn(pred, y).item()
    test_loss /= num_batches
    print(f'test_loss:{test_loss}')

for t in range(epochs):
    print(f'epoch {t+1}\n---------------------------')
    train(train_dataloader, model, loss_fn, optimizer, device)
    test(test_dataloader, model, loss_fn)