본문 바로가기

카테고리 없음

torch dataloader 사용법

728x90
반응형
import torch
from torch import nn, optim
from sklearn.datasets import load_iris
from torch.utils.data import  TensorDataset, DataLoader
iris = load_iris()
X = iris.data[:100]
y = iris.target[:100]
X = torch.tensor(X, dtype=torch.float32)
y = torch.tensor(y, dtype=torch.float32)

# X와 Y를 결합할 때는 아래와 같이 사용.
ds = TensorDataset(X, y)
loader = DataLoader(ds, batch_size=10, shuffle=True)
net = nn.Linear(4, 1)

loss_fn = nn.BCEWithLogitsLoss()
optimizer = optim.SGD(net.parameters(), lr=0.25)
losses = []
for epoc in range(100):
    batch_loss = 0.0
    for xx, yy in loader:
        optimizer.zero_grad()
        y_pred = net(xx)
        loss = loss_fn(y_pred.view_as(yy), yy)
        loss.backward()
        optimizer.step()
        batch_loss += loss.item()
    losses.append(batch_loss)
from matplotlib import pyplot as plt
plt.plot(losses)
plt.show()
728x90
반응형