GeezerVison is showing github.com — Sync2Gen/test_sparse.py at main · ShapeNet/Sync2Gen

Expand file tree

test_sparse.py

Copy path

More file actions

Latest commit

History

277 lines (218 loc) · 11.3 KB

test_sparse.py

Copy path

File metadata and controls

277 lines (218 loc) · 11.3 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

import os

import time

import sys

import shutil

import random

import argparse

import numpy as np

import importlib

import torch

from torch import nn

import torch.utils.data

from dataloader.loader_discrete import DatasetDiscrete

import vis_utils

torch.backends.cudnn.benchmark = True

torch.backends.cudnn.deterministic = True

from scene_loss_utils import get_loss

sys.path.append('models')

from tf_visualizer import Visualizer as TfVisualizer

parser = argparse.ArgumentParser()

# data parameters

parser.add_argument('--data_path', type=str, default='/mnt/yanghaitao/Dataset/Scene_Dataset/3D-FRONT/src/dataset/data/')

parser.add_argument('--type', type=str, help='bedroom or living')

parser.add_argument('--checkpoint_path', default=None, help='Model checkpoint path [default: None]')

parser.add_argument('--log_dir', default='log', help='Dump dir to save model checkpoint [default: log]')

parser.add_argument('--dump_dir', default='dump', help='Dump dir to save results')

parser.add_argument('--not_load_model', action='store_false', help='whether load checkpoint')

# training paramters

parser.add_argument('--max_epoch', type=int, default=1000, help='Epoch to run [default: 180]')

parser.add_argument('--batch_size', type=int, default=32, help='Batch Size during training [default: 8]')

parser.add_argument('--learning_rate', type=float, default=0.001, help='Initial learning rate [default: 0.001]')

parser.add_argument('--weight_decay', type=float, default=0, help='Optimization L2 weight decay [default: 0]')

parser.add_argument('--bn_decay_step', type=int, default=20, help='Period of BN decay (in epochs) [default: 20]')

parser.add_argument('--bn_decay_rate', type=float, default=0.5, help='Decay rate for BN decay [default: 0.5]')

parser.add_argument('--lr_decay_steps', default='400,700,1000', help='When to decay the learning rate (in epochs) [default: 80,120,160]')

parser.add_argument('--lr_decay_rates', default='0.1,0.1,0.1', help='Decay rates for lr decay [default: 0.1,0.1,0.1]')

# model parameters

parser.add_argument('--model_dict', type=str, default='model_scene_box_conv_sparse', help='model file name')

parser.add_argument('--variational', action='store_true', help='whether use VAE model')

parser.add_argument('--max_parts', type=int, default=120, help='maximum number of parts in one object')

parser.add_argument('--num_class', type=int, default=30, help='number of semantic class')

parser.add_argument('--num_each_class', type=int, default=4, help='number of each semantic class')

parser.add_argument('--adjust_kld', action='store_true', help='whether adjust kld for different epochs')

parser.add_argument('--weight_kld', type=float, default=0.001, help='weight of kldiv loss')

parser.add_argument('--weight_rep', type=float, default=1.0, help='weight of representation loss')

parser.add_argument('--kld_interval', type=int, default=50, help='interval of kldiv loss cycle')

parser.add_argument('--latent_dim', type=int, default=256, help='latent vector dimension')

parser.add_argument('--abs_dim', type=int, default=16, help='abs dimension')

parser.add_argument('--sparse_num', type=int, default=4, help='sparse linear layer parameter, switch to fully linear if 0')

parser.add_argument('--valid_threshold', type=float, default=0.5, help='output valid mask threshold')

parser.add_argument('--use_dumped_pairs', action='store_true', help='use the dumped pairs')

# general parameters

parser.add_argument('--vis', action='store_true', help='whether do the visualization')

parser.add_argument('--eval', action='store_true', help='whether switch to eval module')

parser.add_argument('--debug', action='store_true', help='whether switch to debug module')

parser.add_argument('--dump_results', action='store_true', help='whether dump predicted result')

parser.add_argument('--eval_epoch', type=int, default=799, help='Epoch to eval')

parser.add_argument('--gen_from_noise', action='store_true', help='whether generate from noise')

parser.add_argument('--num_gen_from_noise', type=int, default=100, help='num of generation samples')

args = parser.parse_args()

MAX_EPOCH = args.max_epoch

BATCH_SIZE = args.batch_size

BASE_LEARNING_RATE = args.learning_rate

BN_DECAY_STEP = args.bn_decay_step

BN_DECAY_RATE = args.bn_decay_rate

LR_DECAY_STEPS = [int(x) for x in args.lr_decay_steps.split(',')]

LR_DECAY_RATES = [float(x) for x in args.lr_decay_rates.split(',')]

assert(len(LR_DECAY_STEPS)==len(LR_DECAY_RATES))

if args.adjust_kld:

assert(args.variational)

assert(args.type == 'bedroom' or args.type == 'living')

DATA_PATH = os.path.join(args.data_path, args.type)

TRAIN_DATASET = f'{args.data_path}/train_{args.type}.txt'

VAL_DATASET = f'{args.data_path}/val_{args.type}.txt'

# Prepare LOG_DIR

LOG_DIR = args.log_dir

MODEL_NAME = LOG_DIR.split('/')[-1]

if not os.path.exists(LOG_DIR):

os.makedirs(LOG_DIR)

DEFAULT_CHECKPOINT_PATH = os.path.join(LOG_DIR, f'checkpoint_eval{args.eval_epoch}.tar')

CHECKPOINT_PATH = args.checkpoint_path if args.checkpoint_path is not None \

else DEFAULT_CHECKPOINT_PATH

# Init datasets and dataloaders

def my_worker_init_fn(worker_id):

np.random.seed(np.random.get_state()[1][0] + worker_id)

# Create dataset

train_dataset = DatasetDiscrete(DATA_PATH, TRAIN_DATASET)

val_dataset = DatasetDiscrete(DATA_PATH, VAL_DATASET)

num_workers = 0 if args.debug else 4

val_dataloader = torch.utils.data.DataLoader(val_dataset, batch_size=BATCH_SIZE, \

shuffle=False, num_workers=num_workers, worker_init_fn=my_worker_init_fn)

# Initialize the model and optimizer

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

importlib.invalidate_caches()

model_dict = importlib.import_module(args.model_dict)

if args.variational:

net = model_dict.VAE(abs_dim=args.abs_dim, num_class=args.num_class, num_each_class=args.num_each_class, latent_dim=args.latent_dim,

use_dumped_pairs=args.use_dumped_pairs, variational=args.variational, log_dir=LOG_DIR)

else:

raise AssertionError('AE is currently not checked. Need to check model')

print(net)

if torch.cuda.device_count() > 1:

print("Let's use %d GPUs!" % (torch.cuda.device_count()))

net = nn.DataParallel(net)

net.to(device)

# Load the Adam optimizer

optimizer = torch.optim.Adam(net.parameters(), lr=BASE_LEARNING_RATE, weight_decay=args.weight_decay)

# Load checkpoint if any

start_epoch = 0

if CHECKPOINT_PATH is not None and os.path.isfile(CHECKPOINT_PATH) and args.not_load_model:

print('load checkpoint path: %s' % CHECKPOINT_PATH)

checkpoint = torch.load(CHECKPOINT_PATH)

model_state_dict_new = {}

for key, val in checkpoint['model_state_dict'].items():

if 'pairs' not in key:

model_state_dict_new[key] = val

net.load_state_dict(model_state_dict_new)

optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

start_epoch = checkpoint['epoch']

print("Successfully Load Model...")

# Helper for learning rate adjustment

def get_current_lr(epoch):

lr = BASE_LEARNING_RATE

for i,lr_decay_epoch in enumerate(LR_DECAY_STEPS):

if epoch >= lr_decay_epoch:

lr *= LR_DECAY_RATES[i]

return lr

def adjust_learning_rate(optimizer, epoch):

lr = get_current_lr(epoch)

for param_group in optimizer.param_groups:

param_group['lr'] = lr

def adjust_weight_kld(epoch):

if epoch < args.kld_interval:

return (epoch / args.kld_interval) * args.weight_kld

else:

return args.weight_kld

# TFBoard visualizer

TRAIN_VISUALIZER = TfVisualizer(LOG_DIR, 'train')

VAL_VISUALIZER = TfVisualizer(LOG_DIR, 'val')

def eval_one_epoch(epoch, eval_mode):

assert eval_mode in ['val']

print('================ In ' + eval_mode + ' mode ================')

if args.adjust_kld:

weight_kld = adjust_weight_kld(epoch)

else:

weight_kld = args.weight_kld

stat_dict = {}

net.eval()

eval_dataloader = val_dataloader

if args.dump_results:

DUMP_DIR = f'{args.dump_dir}/{args.type}/{MODEL_NAME}'

if not os.path.exists(DUMP_DIR):

os.makedirs(DUMP_DIR)

'''

Note that we set torch.manual_seed(0)

'''

if args.gen_from_noise:

assert(args.variational)

num_epochs = args.num_gen_from_noise // BATCH_SIZE

num_rest = args.num_gen_from_noise % BATCH_SIZE

count = 0

for i in range(num_epochs + 1):

print(i)

with torch.no_grad():

if i < num_epochs:

num_noise = BATCH_SIZE

else:

num_noise = num_rest

latent_code = torch.randn(num_noise, args.latent_dim).to(device)

pred_abs, pred_rel = net(None, None, latent_code=latent_code)

if args.dump_results:

for idx in range(num_noise):

np.save(os.path.join(DUMP_DIR, str(count).zfill(4) + '_abs_pred.npy'), pred_abs[idx].detach().cpu().numpy())

np.save(os.path.join(DUMP_DIR, str(count).zfill(4) + '_rel_pred.npy'), pred_rel[idx].detach().cpu().numpy())

count += 1

assert(count == args.num_gen_from_noise)

else:

for batch_idx, batch_data_label in enumerate(eval_dataloader):

if batch_idx % 10 == 0:

print('Eval batch: %d'%(batch_idx))

for key in batch_data_label:

batch_data_label[key] = batch_data_label[key].to(device)

# Forward pass

inputs_abs = batch_data_label['X_abs']

inputs_rel = batch_data_label['X_rel']

labels_abs = batch_data_label['X_abs']

labels_rel = batch_data_label['X_rel']

with torch.no_grad():

pred_abs, pred_rel, kldiv_loss = net(inputs_abs, inputs_rel)

total_loss, loss_dict, _, _, _ = get_loss(args.type, labels_abs, pred_abs, labels_rel, pred_rel,

num_class=args.num_class,

num_each_class=args.num_each_class,

)

if args.variational:

total_loss += weight_kld * kldiv_loss

loss_dict['kldiv_loss'] = weight_kld * kldiv_loss

loss_dict['loss'] += weight_kld * kldiv_loss

# Accumulate statistics and print out

for key in loss_dict:

if key not in stat_dict: stat_dict[key] = 0

stat_dict[key] += loss_dict[key].item()

if args.dump_results:

if not os.path.exists(f"joint_hyper_opt/{DUMP_DIR}"):

os.makedirs(f"joint_hyper_opt/{DUMP_DIR}")

for b in range(batch_data_label['X_abs'].shape[0]):

fid = str(batch_data_label['index'][b].cpu().numpy()).zfill(4)

print(fid)

np.save(f'joint_hyper_opt/{DUMP_DIR}/{fid}_abs_gt.npy', labels_abs[b].detach().cpu().numpy())

np.save(f'joint_hyper_opt/{DUMP_DIR}/{fid}_abs_pred.npy', pred_abs[b].detach().cpu().numpy())

np.save(f'joint_hyper_opt/{DUMP_DIR}/{fid}_rel_pred.npy', pred_rel[b].detach().cpu().numpy())

for key in sorted(stat_dict.keys()):

print('%s %s: %f' % (eval_mode, key, stat_dict[key] / (float(batch_idx+1))), end=' ')

return 0

def test(eval_epoch):

np.random.seed(0)

torch.manual_seed(0)

eval_one_epoch(eval_epoch, 'val')

if __name__ == '__main__':

test(args.eval_epoch)