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

Expand file tree

train_sparse.py

Copy path

More file actions

Latest commit

History

327 lines (251 loc) · 12.7 KB

train_sparse.py

Copy path

File metadata and controls

327 lines (251 loc) · 12.7 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

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

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('--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')

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, 'checkpoint.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

train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, \

shuffle=True, num_workers=num_workers, worker_init_fn=my_worker_init_fn)

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)

net.load_state_dict(checkpoint['model_state_dict'])

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 train_one_epoch(epoch):

stat_dict = {}

adjust_learning_rate(optimizer, epoch)

if args.adjust_kld:

weight_kld = adjust_weight_kld(epoch)

else:

weight_kld = args.weight_kld

print('Current weight of kld loss: %f'% weight_kld)

net.train()

for batch_idx, batch_data_label in enumerate(train_dataloader):

for key in batch_data_label:

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

# Forward pass

optimizer.zero_grad()

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']

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

total_loss, loss_dict, batch_index, matched_gt_idx, matched_pred_idx = 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

total_loss.backward()

optimizer.step()

# 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()

batch_interval = 50

if (batch_idx+1) % batch_interval == 0:

print('batch: %03d:' % (batch_idx+1), end=' ')

TRAIN_VISUALIZER.log_scalars({key: stat_dict[key] / batch_interval for key in stat_dict},

(epoch*len(train_dataloader)+batch_idx)*BATCH_SIZE)

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

print('%s: %f |' % (key, stat_dict[key] / batch_interval), end=' ')

stat_dict[key] = 0

print()

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

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():

if args.variational:

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

total_loss, loss_dict, batch_index, matched_gt_idx, matched_pred_idx = 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()

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

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

mean_loss = stat_dict['loss']/float(batch_idx+1)

if args.debug:

import ipdb; ipdb.set_trace()

return mean_loss

if args.vis or args.eval or args.dump_results:

return mean_loss

# Log statistics

VAL_VISUALIZER.log_scalars({key: stat_dict[key] / float(batch_idx+1) for key in stat_dict},

(epoch+1) * len(train_dataloader) * BATCH_SIZE)

print()

return mean_loss

def train(start_epoch):

for epoch in range(start_epoch, MAX_EPOCH):

print('**** EPOCH %03d ****' % (epoch))

print('Current learning rate: %f'%(get_current_lr(epoch)))

if args.vis or args.eval:

val_loss = eval_one_epoch(epoch, 'val')

break

# Reset numpy seed.

# REF: https://github.com/pytorch/pytorch/issues/5059

''' here: problem: whether need to use seed '''

np.random.seed()

train_one_epoch(epoch)

# Eval every 10 epochs

if epoch == 0 or epoch % 10 == 9:

val_loss = eval_one_epoch(epoch, 'val')

# Save checkpoint

save_dict = {'epoch': epoch+1, # after training one epoch, the start_epoch should be epoch+1

'optimizer_state_dict': optimizer.state_dict(),

'loss': val_loss,

}

try: # with nn.DataParallel() the net is added as a submodule of DataParallel

save_dict['model_state_dict'] = net.module.state_dict()

except:

save_dict['model_state_dict'] = net.state_dict()

torch.save(save_dict, os.path.join(LOG_DIR, 'checkpoint.tar'))

if epoch > 0 and epoch % 100 == 99:

torch.save(save_dict, os.path.join(LOG_DIR, 'checkpoint_eval%d.tar' % epoch))

if __name__ == '__main__':

train(start_epoch)