小编饼干大*_*干大仙的帖子

我刚刚开始学习 Pytorch 并创建了我的第一个 CNN。该数据集包含 3360 张 RGB 图像，我将它们转换为[3360, 3, 224, 224]张量。数据和标签在dataset(torch.utils.data.TensorDataset). 下面是训练代码。

    def train_net():
        dataset = ld.load()
        data_iter = Data.DataLoader(dataset, batch_size=168, shuffle=True)
        net = model.VGG_19()
        summary(net, (3, 224, 224), device="cpu")
        loss_func = nn.CrossEntropyLoss()
        optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9, dampening=0.1)
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)
        for epoch in range(5):
            print("epoch:", epoch + 1)
            train_loss = 0
            for i, data in enumerate(data_iter, 0):
                x, y = data
                print(x.dtype)
                optimizer.zero_grad()
                out = net(x)
                loss = loss_func(out, y)
                loss.backward()
                optimizer.step() …