def train_with_data_aug(train_augs, test_augs, lr=0.001):
batch_size, net = 2, d2l.resnet18(10)
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
loss = torch.nn.CrossEntropyLoss()
train_iter = load_cifar10(True, train_augs, batch_size)
test_iter = load_cifar10(False, test_augs, batch_size)
train(train_iter, test_iter, net, loss, optimizer, device, num_epochs=10)
def train_with_data_aug(train_augs, test_augs, lr=0.001):
# 数据
batch_size = 256
train_iter = load_cifar10(True, train_augs, batch_size) # 训练数据的迭代器
test_iter = load_cifar10(False, test_augs, batch_size) # 预测数据的迭代器
net = d2l.resnet18(10) # 模型，输出10种分类概率
num_epochs = 10 # 参数
optimizer = torch.optim.Adam(net.parameters(), lr=lr) # 优化器
loss = torch.nn.CrossEntropyLoss() # 损失函数
d2l.train(train_iter, test_iter, net, loss, optimizer, device,
num_epochs) # 训练