Chapter 03無料公開

Auto Encoderの学習

(株)ディープブラック
(株)ディープブラック
2021.02.01に更新

ライブラリのimport

import torch
import torch.utils.data
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
import pickle
from detection_model import autoencoder # 有料

データのロード

with open('normal_data.pkl', 'rb') as f:
	normal_data = pickle.load(f)
with open('anomaly_data.pkl', 'rb') as f:
	anomaly_data = pickle.load(f)

ハイパーパラメータ

Deep Learningのハイパーパラメータを下記のように設定します。

# ハイパーパラメータ
LEARNING_RATE = 0.0001
BATCH_SIZE = 10
EPOCHS = 50

Auto Encoderの訓練データの設定

Auto Encoderで訓練するデータを設定します。

train_data = normal_data[:900]

モデルの設定

model = autoencoder()

DataSet、DataLoaderの設定

class MyDataSet(torch.utils.data.Dataset):
	def __init__(self, data):

		self.data = data
		self.length = len(data)

	def __len__(self):

		return self.length

	def __getitem__(self, index):

		data = self.data[index]

		return data
		
trainset = MyDataSet(train_data)
trainloader = torch.utils.data.DataLoader(
	trainset, batch_size=BATCH_SIZE, shuffle=True)

損失関数、最適化手法の設定

criterion = nn.MSELoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

学習

for epoch in range(EPOCHS):
	print(epoch)
	train_loss = 0
	total = 0
	model.train()
	# cnt = 0
	for data in trainloader:

		optimizer.zero_grad()

		output = model(data[0].float())
		target  = data[0].float()

		loss = criterion(output, target)

		train_loss += loss.item()

		total += data[1].size(0)

		loss.backward()
		optimizer.step()

	train_loss = train_loss / total
	print(f"{train_loss}")

モデルの保存

torch.save(model.state_dict(), "autoencoder.pth")