import sys import os import torch import numpy as np import pandas as pd import plotly.express as px from torchvision.utils import save_image from model import Autoencoder from model import Generator from model import Discriminator from model import AttentiveRNN from tools.argparser import get_param from tools.logger import Logger from tools.dataloader import Dataset # this function is from https://github.com/eriklindernoren/PyTorch-GAN/blob/master/implementations/dualgan/models.py def weights_init_normal(m): classname = m.__class__.__name__ if classname.find("Conv") != -1: torch.nn.init.normal_(m.weight.data, 0.0, 0.02) elif classname.find("BatchNorm2d") != -1: torch.nn.init.normal_(m.weight.data, 1.0, 0.02) torch.nn.init.constant_(m.bias.data, 0.0) param = get_param() logger = Logger() cuda = True if torch.cuda.is_available() else False generator = Generator() # get network values and stuff discriminator = Discriminator() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') generator = Generator().to(device) discriminator = Discriminator().to(device) if load is not False: generator.load_state_dict(torch.load("example_path")) discriminator.load_state_dict(torch.load("example_path")) else: generator.apply(weights_init_normal) discriminator.apply(weights_init_normal) dataloader = Dataloader() # declare generator loss for epoch in range(epoch): for i, (imgs, _) in enumerate(dataloader): torch.save(generator.attentionRNN.state_dict(), "attentionRNN_model_path") ## RNN 따로 돌리고 CPU로 메모리 옳기고 ## Autoencoder 따로 돌리고 메모리 옳기고 ## 안되는가 ## 대충 열심히 GAN 구성하는 코드 ## 대충 그래서 weight export해서 inference용과 training용으로 나누는 코드 ## 대충 그래서 inference용은 attention map까지 하는 녀석과 deraining까지 하는 녀석 두개가 나오는 코드 ## 학습용은 그래서 풀 weight 나옴 ## GAN은 학습 시키면 Nash equilibrium ... 나오게 할 수 있으려나? ## 대충 학습은 어떻게 돌려야 되지 하는 코드 ## generator에서 튀어 나온 애들을 따로 저장해야 하는건가