라이브러리 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)

하이퍼파라미터

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

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")

계속