Spaces:
Sleeping
Sleeping
File size: 3,088 Bytes
0917e8d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import json
import torch
import os
class Config:
"""Config class"""
def __init__(self, tag, root=""):
self.tag = tag
self.cli = False
# self.wandb = True
self.path = os.path.join(root, f"runs/{self.tag}")
self.cm = "gray"
self.data_path = ""
self.mask_coords = []
self.net_type = "conv-resize"
self.image_type = "n-phase"
self.l = 80
self.n_phases = 2
# Training hyperparams
self.batch_size = 4
self.beta1 = 0.9
self.beta2 = 0.999
self.max_iters = 400e3
self.timeout = 1e12
self.lrg = 0.0005
self.lr = 0.0005
self.Lambda = 10
self.critic_iters = 10
self.pw_coeff = 1
self.ngpu = torch.cuda.device_count()
if self.ngpu > 0:
self.device_name = "cuda:0"
else:
self.device_name = "cpu"
self.conv_resize = True
self.nz = 100
# Architecture
self.lays = 4
self.laysd = 5
# kernel sizes
self.dk, self.gk = [4] * self.laysd, [4] * self.lays
self.ds, self.gs = [2] * self.laysd, [2] * self.lays
self.df, self.gf = [self.n_phases, 64, 128, 256, 512, 1], [
self.nz,
512,
256,
128,
self.n_phases,
]
self.dp, self.gp = [1] * self.laysd, [2] * self.lays
# Last two layers conv resize (3,1,0)
self.gk[-2:], self.gs[-2:], self.gp[-2:] = [3, 3], [1, 1], [0, 0]
def update_params(self):
self.df[0] = self.n_phases
self.gf[-1] = self.n_phases
def save(self):
# j = {}
# for k, v in self.__dict__.items():
# j[k] = v
# with open(f"{self.path}/config.json", "w") as f:
# json.dump(j, f)
pass
def load(self):
with open(f"{self.path}/config.json", "r") as f:
j = json.load(f)
for k, v in j.items():
setattr(self, k, v)
def get_net_params(self):
return self.dk, self.ds, self.df, self.dp, self.gk, self.gs, self.gf, self.gp
def get_train_params(self):
return (
self.l,
self.batch_size,
self.beta1,
self.beta2,
self.lrg,
self.lr,
self.Lambda,
self.critic_iters,
self.nz,
)
class ConfigPoly(Config):
def __init__(self, tag, root):
super(ConfigPoly, self).__init__(tag, root=root)
self.frames = 100
# optimisation parameters
if self.cli:
self.opt_iters = 10000
else:
self.opt_iters = 1000
self.opt_lr = 0.001
# if self.image_type=='colour':
self.opt_kl_coeff = 0.00001
def get_train_params(self):
return (
self.l,
self.batch_size,
self.beta1,
self.beta2,
self.lrg,
self.lr,
self.Lambda,
self.critic_iters,
self.nz,
)
|