# テスト画像のバッチの予測をログに記録するための便利な関数
def log_test_predictions(images, labels, outputs, predicted, test_table, log_counter):
# すべてのクラスの信頼度スコアを取得
scores = F.softmax(outputs.data, dim=1)
log_scores = scores.cpu().numpy()
log_images = images.cpu().numpy()
log_labels = labels.cpu().numpy()
log_preds = predicted.cpu().numpy()
# 画像の順序に基づいた ID の追加
_id = 0
for i, l, p, s in zip(log_images, log_labels, log_preds, log_scores):
# データテーブルに必要な情報を追加:
# id, 画像ピクセル, モデルの推測, 正解ラベル, すべてのクラスのスコア
img_id = str(_id) + "_" + str(log_counter)
test_table.add_data(img_id, wandb.Image(i), p, l, *s)
_id += 1
if _id == NUM_IMAGES_PER_BATCH:
break
# W&B: このモデルのトレーニングを追跡するための新しい run を初期化
with wandb.init(project="table-quickstart") as run:
# W&B: config を使用してハイパーパラメーターをログに記録
cfg = run.config
cfg.update({"epochs" : EPOCHS, "batch_size": BATCH_SIZE, "lr" : LEARNING_RATE,
"l1_size" : L1_SIZE, "l2_size": L2_SIZE,
"conv_kernel" : CONV_KERNEL_SIZE,
"img_count" : min(10000, NUM_IMAGES_PER_BATCH*NUM_BATCHES_TO_LOG)})
# モデル、損失関数、オプティマイザーの定義
model = ConvNet(NUM_CLASSES).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
# モデルのトレーニング
total_step = len(train_loader)
for epoch in range(EPOCHS):
# トレーニングステップ
for i, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
# forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# 逆伝播と最適化
optimizer.zero_grad()
loss.backward()
optimizer.step()
# W&B: トレーニングステップごとの損失をログに記録し、UI でライブ可視化
run.log({"loss" : loss})
if (i+1) % 100 == 0:
print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
.format(epoch+1, EPOCHS, i+1, total_step, loss.item()))
# W&B: 各テストステップの予測を保存する Table を作成
columns=["id", "image", "guess", "truth"]
for digit in range(10):
columns.append("score_" + str(digit))
test_table = wandb.Table(columns=columns)
# モデルのテスト
model.eval()
log_counter = 0
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
if log_counter < NUM_BATCHES_TO_LOG:
log_test_predictions(images, labels, outputs, predicted, test_table, log_counter)
log_counter += 1
total += labels.size(0)
correct += (predicted == labels).sum().item()
acc = 100 * correct / total
# W&B: トレーニングエポックごとの精度をログに記録し、UI で可視化
run.log({"epoch" : epoch, "acc" : acc})
print('Test Accuracy of the model on the 10000 test images: {} %'.format(acc))
# W&B: 予測テーブルを wandb にログ記録
run.log({"test_predictions" : test_table})