Pytorch之数据处理

from torch.utils.data import TensorDataset

from torch.utils.data import DataLoader

train_ds = TensorDataset(x_train, y_train)

train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)

valid_ds = TensorDataset(x_valid, y_valid)

valid_dl = DataLoader(valid_ds, batch_size=bs * 2)

def get_data(train_ds, valid_ds, bs):

    return (

        DataLoader(train_ds, batch_size=bs, shuffle=True),

        DataLoader(valid_ds, batch_size=bs * 2),

import numpy as np

def fit(steps, model, loss_func, opt, train_dl, valid_dl):

    for step in range(steps):

        model.train()

        for xb, yb in train_dl:

            loss_batch(model, loss_func, xb, yb, opt)

        model.eval()

        with torch.no_grad():

            losses, nums = zip(

                *[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]

        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)

        print('当前step:'+str(step), '验证集损失：'+str(val_loss))

from torch import optim

def get_model():

    model = Mnist_NN()

    return model, optim.SGD(model.parameters(), lr=0.001)

def loss_batch(model, loss_func, xb, yb, opt=None):

    loss = loss_func(model(xb), yb)

    if opt is not None:

        loss.backward()

        opt.step()

        opt.zero_grad()

    return loss.item(), len(xb)

train_dl, valid_dl = get_data(train_ds, valid_ds, bs)

model, opt = get_model()

fit(25, model, loss_func, opt, train_dl, valid_dl)

巴特西