03_computer_vision_dnn (original) (raw)

Get cpu or gpu device for training.

device = "cuda" if torch.cuda.is_available() else "cpu" print("Using {} device".format(device))

def train(dataloader, model, loss_fn, epochs, **kwargs): # initialize model with **kwargs if isinstance(model, type): model = model(**kwargs) optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) size = len(dataloader.dataset) for t in range(epochs): # print(f"Epoch {t+1}\n-------------------------------") for batch, (X, y) in enumerate(dataloader): X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            # print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
return model

def test(dataloader, model, loss_fn): size = len(dataloader.dataset) num_batches = len(dataloader) model.eval() test_loss, correct = 0, 0 with torch.no_grad(): for X, y in dataloader: X, y = X.to(device), y.to(device) pred = model(X) test_loss += loss_fn(pred, y).item() correct += (pred.argmax(1) == y).type(torch.float).sum().item() test_loss /= num_batches correct /= size return correct