def do_detect(model, img, conf_thresh, nms_thresh, use_cuda=1): model.eval() t0 = time.time() if type(img) == np.ndarray and len(img.shape) == 3: # cv2 image img = nd.from_numpy(img.transpose(2, 0, 1)).\ img = nd.broadcast_div(img,255.0) img.expend_dims(axis=0) elif type(img) == np.ndarray and len(img.shape) == 4: img = nd.from_numpy(img.transpose(0, 3, 1, 2)) img = nd.broadcast_div(img,255.0) else: print("unknow image type") exit(-1) if use_cuda: img = img.cuda() img = nd.autograd.Variable(img) t1 = time.time() output = model(img) t2 = time.time() print('-----------------------------------') print(' Preprocess : %f' % (t1 - t0)) print(' Model Inference : %f' % (t2 - t1)) print('-----------------------------------') return utils.post_processing(img, conf_thresh, nms_thresh, output)
def forward(self, outputs, targets_heatmaps, targets_scale, targets_offset, targets_inds, targets_reg_mask): opt = self.opt hm_loss, wh_loss, off_loss = 0, 0, 0 for s in range(opt.num_stacks): output = outputs[s] if not opt.mse_loss: output['hm'] = _sigmoid(output['hm']) # Optional: Use ground truth for validation if opt.eval_oracle_hm: output['hm'] = targets_heatmaps if opt.eval_oracle_wh: output['wh'] = nd.from_numpy( gen_oracle_map(targets_scale.asnumpy(), targets_inds.asnumpy(), output['wh'].shape[3], output['wh'].shape[2])).as_in_context( opt.device) if opt.eval_oracle_offset: output['reg'] = nd.from_numpy( gen_oracle_map(targets_offset.asnumpy(), targets_inds.asnumpy(), output['reg'].shape[3], output['reg'].shape[2])).as_in_context( opt.device) # 1. heatmap loss hm_loss = hm_loss + self.crit(output['hm'], targets_heatmaps) / opt.num_stacks # 2. scale loss if opt.wh_weight > 0: wh_loss = wh_loss + self.crit_reg( output['wh'], targets_reg_mask, targets_inds, targets_scale) / opt.num_stacks # 3. offset loss if opt.reg_offset and opt.off_weight > 0: off_loss = off_loss + self.crit_reg( output['reg'], targets_reg_mask, targets_inds, targets_offset) / opt.num_stacks # total loss loss = opt.hm_weight * hm_loss + opt.wh_weight * wh_loss + \ opt.off_weight * off_loss loss_stats = { 'loss': loss, 'hm_loss': hm_loss, 'wh_loss': wh_loss, 'off_loss': off_loss } return loss
def validate(epoch, val_loader, model, crit_cls, crit_reg, opt, ctx, gen_shape=False): """ One validation """ generated_shapes = [] original_shapes = [] sample_prob = opt.inner_sample_prob loss_cls_sum, loss_reg_sum, n = 0.0, 0.0, 0 for idx, data in enumerate(val_loader): start = time.time() shapes, labels, masks, params, param_masks = data[0], data[1], data[ 2], data[3], data[4] gt = shapes shapes = nd.expand_dims(shapes, axis=1) shapes = shapes.as_in_context(ctx) labels = labels.as_in_context(ctx) masks = masks.as_in_context(ctx) params = params.as_in_context(ctx) param_masks = param_masks.as_in_context(ctx) with autograd.train_mode(): out = model.decode(shapes) #out = model(shapes, labels, sample_prob) bsz, n_block, n_step = labels.shape labels = labels.reshape(bsz, n_block * n_step) masks = masks.reshape(bsz, n_block * n_step) out_pgm = out[0].reshape(bsz, n_block * n_step, opt.program_size + 1) bsz, n_block, n_step, n_param = params.shape params = params.reshape(bsz, n_block * n_step, n_param) param_masks = param_masks.reshape(bsz, n_block * n_step, n_param) out_param = out[1].reshape(bsz, n_block * n_step, n_param) loss_cls, acc = crit_cls(out_pgm, labels, masks) loss_reg = crit_reg(out_param, params, param_masks) end = time.time() loss_cls = loss_cls.mean().asscalar() loss_reg = loss_reg.mean().asscalar() if idx % opt.info_interval == 0: out_1 = nd.round(out[0]).astype('int64') out_2 = nd.round(out[1]).astype('int64') pred = nd.from_numpy(decode_multiple_block( out_1, out_2)).astype("float32").as_in_context(mx.cpu()) IoU = BatchIoU(pred, gt) print( "Test: epoch {} batch {}/{}, loss_cls = {:.3f}, loss_reg = {:.3f}, acc = {:.3f}, IoU = {:.3f} time = {:.3f}" .format(epoch, idx, len(val_loader), loss_cls, loss_reg, acc[0].asscalar(), IoU.mean(), end - start)) sys.stdout.flush()
def train_step(model, optimizer, data, epoch): running_loss, update_count = 0.0, 0 N = data.shape[0] idxlist = list(range(N)) np.random.shuffle(idxlist) training_steps = len(range(0, N, args.batch_size)) with trange(train_steps) as t: for batch_idx, start_idx in zip(t, range(0, N, args.batch_size)): t.set_description("epoch: {}".format(epoch + 1)) end_idx = min(start_idx + args.batch_size, N) X_inp = data[idxlist[start_idx:end_idx]] X_inp = nd.from_numpy(X_inp.toarray()).as_in_context(ctx) with autograd.record(): if model.__class__.__name__ == "MultiVAE": if args.total_anneal_steps > 0: anneal = min( args.anneal_cap, 1.0 * update_count / args.total_anneal_steps ) else: anneal = args.anneal_cap update_count += 1 loss = model(X_inp, anneal) elif model.__class__.__name__ == "MultiDAE": loss = model(X_inp) trainer.step(X_inp.shape[0]) running_loss += nd.mean(loss).asscalar() avg_loss = running_loss / (batch_idx + 1) t.set_postfix(loss=avg_loss)
def eval_step(data_tr, data_te, data_type="valid"): running_loss, update_count = 0.0, 0 eval_idxlist = list(range(data_tr.shape[0])) eval_N = data_tr.shape[0] eval_steps = len(range(0, eval_N, args.batch_size)) n100_list, r20_list, r50_list = [], [], [] with trange(eval_steps) as t: for batch_idx, start_idx in zip(t, range(0, eval_N, args.batch_size)): t.set_description(data_type) end_idx = min(start_idx + args.batch_size, eval_N) X_tr = data_tr[eval_idxlist[start_idx:end_idx]] X_te = data_te[eval_idxlist[start_idx:end_idx]] X_tr_inp = nd.from_numpy(X_inp.toarray()).as_in_context(ctx) with autograd.predict_mode(): if model.__class__.__name__ == "MultiVAE": if args.total_anneal_steps > 0: anneal = min( args.anneal_cap, 1.0 * update_count / args.total_anneal_steps ) else: anneal = args.anneal_cap loss = model(X_tr_inp, anneal) elif model.__class__.__name__ == "MultiDAE": loss = models(X_tr_inp) running_loss += loss.item() avg_loss = running_loss / (batch_idx + 1) # Exclude examples from training set X_out = X_out.asnumpy() X_out[X_tr.nonzero()] = -np.inf n100 = NDCG_binary_at_k_batch(X_out, X_te, k=100) r20 = Recall_at_k_batch(X_out, X_te, k=20) r50 = Recall_at_k_batch(X_out, X_te, k=50) n100_list.append(n100) r20_list.append(r20) r50_list.append(r20) t.set_postfix(loss=avg_loss) n100_list = np.concatenate(n100_list) r20_list = np.concatenate(r20_list) r50_list = np.concatenate(r50_list) return avg_loss, np.mean(n100_list), np.mean(r20_list), np.mean(r50_list)
def forward(self, graph, feat, etypes): """Compute Gated Graph Convolution layer. Parameters ---------- graph : DGLGraph The graph. feat : mxnet.NDArray The input feature of shape :math:`(N, D_{in})` where :math:`N` is the number of nodes of the graph and :math:`D_{in}` is the input feature size. etypes : torch.LongTensor The edge type tensor of shape :math:`(E,)` where :math:`E` is the number of edges of the graph. Returns ------- mxnet.NDArray The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}` is the output feature size. """ with graph.local_scope(): assert graph.is_homogeneous, \ "not a homogeneous graph; convert it with to_homogeneous " \ "and pass in the edge type as argument" zero_pad = nd.zeros( (feat.shape[0], self._out_feats - feat.shape[1]), ctx=feat.context) feat = nd.concat(feat, zero_pad, dim=-1) for _ in range(self._n_steps): graph.ndata['h'] = feat for i in range(self._n_etypes): eids = (etypes.asnumpy() == i).nonzero()[0] eids = nd.from_numpy(eids, zero_copy=True).as_in_context( feat.context).astype(graph.idtype) if len(eids) > 0: graph.apply_edges( lambda edges: {'W_e*h': self.linears[i](edges.src['h'])}, eids) graph.update_all(fn.copy_e('W_e*h', 'm'), fn.sum('m', 'a')) a = graph.ndata.pop('a') feat = self.gru(a, [feat])[0] return feat
def gather(self, dim, index): """ Gathers values along an axis specified by ``dim``. For a 3-D tensor the output is specified by: out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2 Parameters ---------- dim: The axis along which to index index: A tensor of indices of elements to gather Returns ------- Output Tensor """ idx_xsection_shape = index.shape[:dim] + \ index.shape[dim + 1:] self_xsection_shape = self.shape[:dim] + self.shape[dim + 1:] if idx_xsection_shape != self_xsection_shape: raise ValueError( "Except for dimension " + str(dim) + ", all dimensions of index and self should be the same size") if index.dtype != np.dtype('int_'): raise TypeError("The values of index must be integers") data_swaped = nd.swapaxes(self, 0, dim).asnumpy() index_swaped = nd.swapaxes(index, 0, dim).asnumpy() #print(data_swaped,index_swaped) #print("index_swaped\n",index_swaped,index_swaped.shape,"data_swaped\n",data_swaped,data_swaped.shape,'\n') gathered = nd.from_numpy(np.choose( index_swaped, data_swaped)).as_in_context(d2l.try_gpu()) return nd.swapaxes(gathered, 0, dim)
def _init_bias(self, name, data): print('Init', name, data.shape) data[:] = nd.from_numpy(self.initial_bias)
def combine(x, y): y.setflags(write=1) y = nd.from_numpy(y).as_in_context(self.ctx) y = y.expand_dims(axis=1) return nd.concat(x, y, dim=1)
def train(epoch, train_loader, model,loss, optimizer, opt,ctx,train_loss,train_iou): """ one epoch training for program executor """ loss_sum,iou_sum,n = 0.0,0.0,0 for idx, data in enumerate(train_loader): start_t = time.time() shape, label, param = data bsz = shape.shape[0] n_step = label.shape[1] #print("label.shape:",label) #print("n_step:",n_step,"bsz:",bsz,"stop_id:",stop_id) index = np.array(list(map(lambda x: n_step, label)))-1 #index = label # add noise during training, making the executor accept # continuous output from program generator label = label.reshape(-1,1).asnumpy() pgm_vector = 0.2 * np.random.uniform(0,1,(bsz * n_step, stop_id)) pgm_noise = 0.2 *np.random.uniform(0,1,label.shape) pgm_value = 1 - pgm_noise #print('pgm_val.shape:',pgm_value.shape,'label.shape:',label.shape,'label.shape:',label.shape) pgm_vector = scatter_numpy(pgm_vector,1,label,pgm_value).reshape(bsz,n_step,stop_id) param_noise = nd.random_uniform(0,1,shape=param.shape) param_vector = param + 0.6 * (param_noise - 0.5) #print("param_vector.shape:",param_vector.shape) gt = shape.as_in_context(ctx) #print(pgm_vector.dtype) index = nd.from_numpy(index).astype('int64').as_in_context(ctx) pgm_vector = nd.from_numpy(pgm_vector).astype('float32').as_in_context(ctx) param_vector = param_vector.as_in_context(ctx) with autograd.record(): pred = model(pgm_vector, param_vector, index) scores = nd.log_softmax(pred,axis=1) pred0 = scores[:,0].squeeze()*opt.n_weight pred1 = scores[:,1].squeeze()*opt.p_weight l = -nd.where(gt, pred1, pred0).mean((1,2,3)) #l = -(nd.pick(scores1, gt, axis=1, keepdims=True)*opt.n_weight # +nd.pick(scores2,(1-gt), axis=1, keepdims=True)*opt.p_weight).mean((1,2,3,4)) l.backward() #clip_gradient(optimizer, opt.grad_clip) #optimizer._allreduce_grads(); optimizer.step(l.shape[0],ignore_stale_grad=True) l = l.mean().asscalar() pred = nd.softmax(pred,axis = 1) pred = pred[:, 1, :, :, :] s1 = gt.reshape(-1, 32, 32, 32).astype('float32').as_in_context(mx.cpu()) s2 = pred.squeeze().as_in_context(mx.cpu()) #print(s2.shape) s2 = (s2 > 0.5) batch_iou = BatchIoU(s1, s2) iou = batch_iou.mean() end_t = time.time() loss_sum+=l n+=1 iou_sum+=iou if idx % (opt.info_interval * 10) == 0: print("Train: epoch {} batch {}/{}, loss13 = {:.3f}, iou = {:.3f}, time = {:.3f}" .format(epoch, idx, len(train_loader), l, iou, end_t - start_t)) sys.stdout.flush() train_loss.append(loss_sum/n) train_iou.append(iou_sum/n)
def validate(epoch, val_loader, model, loss, opt, ctx,val_loss,val_iou, gen_shape=False): # load pre-fixed randomization try: rand1 = np.load(opt.rand1) rand2 = np.load(opt.rand2) rand3 = np.load(opt.rand3) except: rand1 = np.random.rand(opt.batch_size * opt.seq_length, stop_id).astype(np.float32) rand2 = np.random.rand(opt.batch_size * opt.seq_length, 1).astype(np.float32) rand3 = np.random.rand(opt.batch_size, opt.seq_length, max_param - 1).astype(np.float32) np.save(opt.rand1, rand1) np.save(opt.rand2, rand2) np.save(opt.rand3, rand3) generated_shapes = None original_shapes = None loss_sum,iou_sum,n = 0.0,0.0,0 for idx, data in enumerate(val_loader): start_t = time.time() shape, label, param = data bsz = shape.shape[0] n_step = label.shape[1] index = np.array(list(map(lambda x: n_step, label))) index = index - 1 # add noise during training, making the executor accept # continuous output from program generator label = label.reshape(-1,1).asnumpy() pgm_vector = 0.1*rand1 pgm_noise = 0.1*rand2 pgm_value = np.ones(label.shape) - pgm_noise #print('pgm_val.shape:',pgm_value.shape,'label.shape:',label.shape,'label.shape:',label.shape) pgm_vector = scatter_numpy(pgm_vector,1,label,pgm_value).reshape(bsz,n_step,stop_id) param_noise = nd.from_numpy(rand3) #print(param.shape,param_noise.shape) param_vector = param + 0.6 * (param_noise - 0.5) gt = shape.astype('float32').as_in_context(ctx) index = nd.from_numpy(index).astype('int64').as_in_context(ctx) pgm_vector = nd.from_numpy(pgm_vector).as_in_context(ctx) param_vector = param_vector.as_in_context(ctx) #prediction pred = model(pgm_vector, param_vector, index) scores = nd.log_softmax(pred,axis=1) pred0 = scores[:,0].squeeze()*opt.p_weight pred1 = scores[:,1].squeeze()*opt.n_weight l = -nd.where(gt, pred1, pred0).mean((1,2,3)) #print(pred2.dtype,gt.dtype) #l = loss(pred,gt,sample_weight = nd.array([opt.n_weight,opt.p_weight])) l = l.mean().asscalar() pred = nd.softmax(pred,axis=1) pred = pred[:, 1, :, :, :] s1 = gt.reshape(-1, 32, 32, 32).as_in_context(mx.cpu()) s2 = pred.squeeze().as_in_context(mx.cpu()) s2 = (s2 > 0.5) batch_iou = BatchIoU(s1, s2) iou = batch_iou.mean() loss_sum+=l n+=1 iou_sum+=iou if(idx+1)%5==0 and gen_shape: if original_shapes is None: original_shapes = s1.expand_dims(axis=0) generated_shapes = s2.expand_dims(axis=0) else: original_shapes = nd.concat(original_shapes,s1.expand_dims(axis=0),dim=0) generated_shapes = nd.concat(generated_shapes,s2.expand_dims(axis=0),dim=0) end_t = time.time() if (idx + 1) % opt.info_interval == 0: print("Test: epoch {} batch {}/{}, loss13 = {:.3f}, iou = {:.3f}, time = {:.3f}" .format(epoch, idx + 1, len(val_loader), l, iou, end_t - start_t)) sys.stdout.flush() if(idx+1>len(val_loader)/10): break; val_loss.append(loss_sum/n) val_iou.append(iou_sum/n) return generated_shapes, original_shapes
def train(epoch, train_loader, generator, executor, criterion, optimizer, opt, ctx): """ one epoch guided adaptation """ def set_bn_eval(m): if m.prefix[:9] == 'batchnorm': m._kwargs['use_global_stats'] = True m.grad_req = 'null' executor.apply(set_bn_eval) for idx, data in enumerate(train_loader): start = time.time() shapes = data.as_in_context(ctx) raw_shapes = data shapes = shapes.expand_dims(axis=1) with autograd.record(): pgms, params = generator.decode(shapes) # truly rendered shapes pgms_int = nd.round(pgms).astype('int64') params_int = nd.round(params).astype('int64') # neurally rendered shapes pgms = nd.exp(pgms) bsz, n_block, n_step, n_vocab = pgms.shape pgm_vector = pgms.reshape(bsz * n_block, n_step, n_vocab) bsz, n_block, n_step, n_param = params.shape param_vector = params.reshape(bsz * n_block, n_step, n_param) index = (n_step - 1) * nd.ones(bsz * n_block).astype('int64') index = index.as_in_context(ctx) pred = executor(pgm_vector, param_vector, index) pred = nd.softmax(pred, axis=1) #print(pred.shape) pred = pred[:, 1] pred = pred.reshape(bsz, n_block, 32, 32, 32) rec = nd.max(pred, axis=1) #print("rec.shape:",rec.shape,"shapes.shape:",shapes.shape) #rec1 = rec.expand_dims(axis=1) rec1 = nd.log(rec + 1e-11) rec0 = nd.log(1 - rec + 1e-11) #rec_all = nd.concat(rec0, rec1, dim=1) #rec_all1 = nd.log(rec_all + 1e-10) #rec_all2 = nd.log(1-rec_all + 1e-10) gt = shapes.squeeze().astype('int64') loss = -nd.where(gt, rec1, rec0).mean(axis=(1, 2, 3)) #loss = -(nd.pick(rec_all1,gt,axis = 1,keepdims=True)).mean(axis = (1,2,3,4)) #loss = criterion(rec_all, gt) loss.backward() optimizer.step(loss.shape[0], ignore_stale_grad=True) l = loss.mean().asscalar() rendered_shapes = decode_multiple_block(pgms_int, params_int) rendered_shapes = nd.from_numpy(rendered_shapes).astype( 'float32').as_in_context(mx.cpu()) IoU2 = BatchIoU(raw_shapes, rendered_shapes) reconstruction = (rec.as_in_context(mx.cpu()) > 0.5).astype('float32') IoU1 = BatchIoU(reconstruction, raw_shapes) #print("IoU1:",IoU1,IoU2) IoU1 = IoU1.mean() IoU2 = IoU2.mean() end = time.time() if idx % opt.info_interval == 0: print( "Train: epoch {} batch {}/{}, loss = {:.3f}, IoU1 = {:.3f}, IoU2 = {:.3f}, time = {:.3f}" .format(epoch, idx, len(train_loader), l, IoU1, IoU2, end - start)) sys.stdout.flush()
def run(): # get options opt = options_guided_adaptation.parse() opt_gen = options_train_generator.parse() opt_exe = options_train_executor.parse() print('===== arguments: guided adaptation =====') for key, val in vars(opt).items(): print("{:20} {}".format(key, val)) print('===== arguments: guided adaptation =====') if not os.path.isdir(opt.save_folder): os.makedirs(opt.save_folder) # build loaders train_set = ShapeNet3D(opt.train_file) train_loader = gdata.DataLoader(dataset=train_set, batch_size=opt.batch_size, shuffle=True, num_workers=opt.num_workers) val_set = ShapeNet3D(opt.val_file) val_loader = gdata.DataLoader(dataset=val_set, batch_size=opt.batch_size, shuffle=False, num_workers=opt.num_workers) def visual(path, epoch, gen_shapes, file_name, nums_samples): data = gen_shapes.transpose((0, 3, 2, 1)) data = np.flip(data, axis=2) num_shapes = data.shape[0] for i in range(min(nums_samples, num_shapes)): voxels = data[i] save_name = os.path.join(path, file_name.format(epoch, i)) visualization(voxels, threshold=0.1, save_name=save_name, uniform_size=0.9) ctx = d2l.try_gpu() # load program generator generator = BlockOuterNet(opt_gen) generator.init_blocks(ctx) generator.load_parameters("model of blockouternet") # load program executor executor = RenderNet(opt_exe) executor.initialize(init=init.Xavier(), ctx=ctx) executor.load_parameters("model of executor") # build loss functions criterion = gloss.SoftmaxCrossEntropyLoss(axis=1, from_logits=True) optimizer = Trainer( generator.collect_params(), "adam", { "learning_rate": opt.learning_rate, "wd": opt.weight_decay, 'beta1': opt.beta1, 'beta2': opt.beta2, 'clip_gradient': opt.grad_clip }) print("###################") print("testing") gen_shapes, ori_shapes = validate(0, val_loader, generator, opt, ctx, gen_shape=True) #visual('imgs of chairs/adaption/chair/',0,ori_shapes,'GT {}-{}.png',8) #visual('imgs of chairs/adaption/chair/',0,gen_shapes,'epoch{}-{}.png',8) gen_shapes = nd.from_numpy(gen_shapes) ori_shapes = nd.from_numpy(ori_shapes) #print(gen_shapes.dtype,ori_shapes.dtype) #print("done",ori_shapes.shape,gen_shapes.shape) IoU = BatchIoU(gen_shapes, ori_shapes) #print(IoU) print("iou: ", IoU.mean()) best_iou = 0 print(opt.epochs) for epoch in range(1, opt.epochs + 1): print("###################") print("adaptation") train(epoch, train_loader, generator, executor, criterion, optimizer, opt, ctx) print("###################") print("testing") gen_shapes, ori_shapes = validate(epoch, val_loader, generator, opt, ctx, gen_shape=True) #visual('imgs of chairs/adaption/chair/',epoch,gen_shapes,'epoch{}-{}.png',8) gen_shapes = nd.from_numpy(gen_shapes) ori_shapes = nd.from_numpy(ori_shapes) IoU = BatchIoU(gen_shapes, ori_shapes) print("iou: ", IoU.mean()) if epoch % opt.save_interval == 0: print('Saving...') generator.save_parameters("generator of GA on shapenet") optimizer.save_states("optimazer of generator of GA on shapenet") if IoU.mean() >= best_iou: print('Saving best model') generator.save_parameters("generator of GA on shapenet") optimizer.save_states("optimazer of generator of GA on shapenet") best_iou = IoU.mean()
def do(self, img, target): if not isinstance(img, nd.NDArray): img = nd.from_numpy(img) img = nd.image.to_tensor(img) target = nd.array(target) return img, target
def train(epoch, train_loader, model, crit_cls, crit_reg, optimizer, opt, ctx): """ One epoch training """ cls_w = opt.cls_weight reg_w = opt.reg_weight # the prob: > 1 # the input of step t Operator where is missing FInferType attributeis always sampled from the output of step t-1 sample_prob = opt.inner_sample_prob for idx, data in enumerate(train_loader): start = time.time() #data, pgm, pgm_mask, param, param_mask shapes, labels, masks, params, param_masks = data[0], data[1], data[ 2], data[3], data[4] gt = shapes shapes = nd.expand_dims(shapes, axis=1) #print(labels[0],params[0]) shapes = shapes.as_in_context(ctx) labels = labels.as_in_context(ctx) labels2 = labels.as_in_context(ctx) masks = masks.as_in_context(ctx) params = params.as_in_context(ctx) param_masks = param_masks.as_in_context(ctx) #shapes.attach_grad(),labels.attach_grad() with autograd.record(): out = model(shapes, labels, sample_prob) #out = model.decode(shapes) # reshape bsz, n_block, n_step = labels.shape labels = labels.reshape(bsz, -1) masks = masks.reshape(bsz, -1) out_pgm = out[0].reshape(bsz, n_block * n_step, opt.program_size + 1) bsz, n_block, n_step, n_param = params.shape params = params.reshape(bsz, n_block * n_step, n_param) param_masks = param_masks.reshape(bsz, n_block * n_step, n_param) out_param = out[1].reshape(bsz, n_block * n_step, n_param) loss_cls, acc = crit_cls(out_pgm, labels, masks) loss_reg = crit_reg(out_param, params, param_masks) loss = cls_w * loss_cls + reg_w * loss_reg loss.backward() optimizer.step(bsz, ignore_stale_grad=True) loss_cls = loss_cls.mean().asscalar() loss_reg = loss_reg.mean().asscalar() end = time.time() if idx % (opt.info_interval * 10) == 0: out_1 = nd.round(out[0]).astype('int64') out_2 = nd.round(out[1]).astype('int64') pred = nd.from_numpy(decode_multiple_block( out_1, out_2)).astype("float32").as_in_context(mx.cpu()) IoU = BatchIoU(pred, gt) print( "Train: epoch {} batch {}/{},loss_cls = {:.3f},loss_reg = {:.3f},acc = {:.3f},IoU = {:.3f},time = {:.2f}" .format(epoch, idx, len(train_loader), loss_cls, loss_reg, acc[0].asscalar(), IoU.mean(), end - start)) sys.stdout.flush()