def train(num_gpus, batch_size, lr):
train_iter, test_iter = util.load_data_fashion_mnist(batch_size)
ctx = [mx.gpu(i) for i in range(num_gpus)]
print('running on:', ctx)
# 将模型参数复制到num_gpus块显卡的显存上
gpu_params = [get_params(params, c) for c in ctx]
for epoch in range(4):
start = time.time()
for X, y in train_iter:
# 对单个小批量进行GPU训练
train_batch(X, y, gpu_params, ctx, lr)
nd.waitall()
train_time = time.time() - start
def net(x): # 在gpu(0)上验证模型
return lenet(x, gpu_params[0])
test_acc = util.evaluate_accuracy(test_iter, net, ctx[0])
print('epoch %d, time %.1f sec, test acc %.2f' %
(epoch + 1, train_time, test_acc))
# 除了最后的卷积层外，进一步增大了输出通道数
# 前两个卷积层后部使用池化层来减小输入的高和宽
nn.Conv2D(384, kernel_size=3, padding=1, activation='relu'),
nn.Conv2D(384, kernel_size=3, padding=1, activation='relu'),
nn.Conv2D(256, kernel_size=3, padding=1, activation='relu'),
nn.MaxPool2D(pool_size=3, strides=2),
# 这里全连接层的输出个数比LeNet中大数倍，使用丢弃层来缓解过拟合
nn.Dense(4096, activation='relu'),
nn.Dropout(0.5),
nn.Dense(4096, activation='relu'),
nn.Dropout(0.5),
# 输出层，对应Fashion-MNIST数据集的10种类别
nn.Dense(10))
X = nd.random.uniform(shape=(1, 1, 224, 224))
net.initialize()
for layer in net:
X = layer(X)
print(layer.name, 'output shape:\t', X.shape)
batch_size = 128
train_iter, test_iter = util.load_data_fashion_mnist(batch_size,
resize=224)
learning_rate = 0.01
num_epochs = 5
ctx = util.try_gpu()
net.initialize(force_reinit=True, ctx=ctx, init=init.Xavier())
trainer = gluon.Trainer(net.collect_params(), 'sgd',
{'learning_rate': learning_rate})
util.train_ch5(net, train_iter, test_iter, batch_size, trainer, ctx,
num_epochs)