# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
import numpy as np

from configuration import Config
from dataset import train_dataset, valid_dataset, test_dataset
from model import DIP

# +
# 하이퍼 파라미터들
C = Config()

valid_sample_num = C.get('valid_sample_num')
test_sample_num = C.get('test_sample_num')

device = C.get('device')
input_size = C.get('input_size')
hidden_size = C.get('hidden_size')
output_size = C.get('output_size')

batch_size = C.get('batch_size')
num_workers = C.get('num_workers')
lr = C.get('lr')
weight_decay = C.get('weight_decay')

epochs = C.get('epochs')
clip_grad = C.get('clip_grad')
print_every = C.get('print_every')
# -

# 데이터 불러오기
train_acc, train_smpl, train_input = train_dataset()
valid_acc, valid_ori, valid_smpl = valid_dataset()
test_acc, test_ori, test_smpl = test_dataset()

# 모델 생성
model = DIP(input_size, hidden_size, output_size, n_layers=2, drop_prob=0.0)
model = model.to(device)
model

criterion = nn.GaussianNLLLoss() # Loss Function
optimizer = torch.optim.Adam(params=model.parameters(), lr=lr)

# train 데이터셋 형성
train_dataset = TensorDataset(torch.from_numpy(train_input),
                              torch.from_numpy(train_smpl),
                              torch.from_numpy(train_acc))
train_loader = DataLoader(train_dataset,
                          batch_size=batch_size,
                          num_workers=num_workers,
                          drop_last=False,
                          shuffle=True)

# +
# validation 데이터셋 형성
valid_input_list, valid_dataset = [], []

for i in range(valid_sample_num):
    valid_input_list.append(torch.concat((torch.tensor([valid_ori[i]]), torch.tensor([valid_acc[i]])), axis=2))
    valid_dataset.append((valid_input_list[i], torch.tensor([valid_smpl[i]]), torch.tensor([valid_acc[i]])))

# +
# test 데이터셋 형성
test_input_list, test_dataset = [], []

for i in range(test_sample_num):
    test_input_list.append(torch.concat((torch.tensor([test_ori[i]]), torch.tensor([test_acc[i]])), axis=2))
    test_dataset.append((test_input_list[i], torch.tensor([test_smpl[i]]), torch.tensor([test_acc[i]])))
# -

if __name__ == "__main__":
    train_losses = []
    valid_losses = []
    accuracies = []
    
    valid_min = np.inf
    
    for e in range(epochs):
        print_count = 0
        
        for train_input, train_smpl, train_acc in train_loader:
            print_count += 1
            
            train_input = train_input.to(device).float()
            train_smpl = train_smpl.to(device).float()
            train_acc = train_acc.to(device).float()
            
            h = model.init_hidden(train_input.shape[0])
            h = tuple([each.data for each in h])
            
            model.zero_grad()
            smpl_mu, acc_mu, smpl_sigma, acc_sigma = model(train_input, h)
            smpl_loss = criterion(smpl_mu, train_smpl, smpl_sigma) # 자세 추정을 위한 Loss
            acc_loss = criterion(acc_mu, train_acc, acc_sigma) # 가속도 복원을 위한 Loss
            final_loss = smpl_loss + acc_loss
            
            final_loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), clip_grad)
            optimizer.step()
            
            if print_count % print_every == 0:
                train_losses.append(final_loss.item())
                
                # Validation 코드 다시 작성해야함
                """
                temp_losses = []
                model.eval()
                
                for valid_input, valid_smpl, valid_acc in valid_dataset:
                    valid_input = valid_input.to(device).float()
                    valid_smpl = valid_smpl.to(device).float()
                    valid_acc = valid_acc.to(device).float()
                    
                    valid_h = model.init_hidden(valid_input.shape[0])
                    valid_h = tuple([each.data for each in valid_h])
                    
                    model.zero_grad()
                    smpl_mu, acc_mu, smpl_sigma, acc_sigma = model(valid_input, valid_h)
                    smpl_loss = criterion(smpl_mu, valid_smpl, smpl_sigma)
                    acc_loss = criterion(acc_mu, valid_acc, acc_sigma)
                    valid_final_loss = smpl_loss + acc_loss
                    
                    temp_losses.append(valid_final_loss.item())
                
                valid_loss = np.mean(temp_losses)
                valid_losses.append(valid_loss)
                
                if valid_loss < valid_min:
                    valid_min = valid_loss
                    torch.save(model.state_dict(), 'model.pt')
                """

                print("Epoch: {}/{}\t".format(e+1, epochs),
                      "Step: {}/{}\t".format(print_count, len(train_loader)),
                      "Train Loss: {:.6f}\t".format(final_loss.item()))
                      #"Valid Loss: {:.6f}".format(valid_loss))
                
            model.train()

# +
import matplotlib.pyplot as plt

plt.plot(np.arange(6000), train_losses)
plt.xlabel("Epochs(x3)")
plt.ylabel("Train Loss")
plt.show()
# -

plt.plot(np.arange(6000), valid_losses)
plt.xlabel("Epochs(x3)")
plt.ylabel("Valid Loss")
plt.show()



StarcoderdataJupyterScriptsDedupFiltered