def execute(self, point_cloud, label): B, D, N = point_cloud.size() trans = self.stn(point_cloud) point_cloud = point_cloud.transpose(0, 2, 1) point_cloud = nn.bmm(point_cloud, trans) point_cloud = point_cloud.transpose(0, 2, 1) out1 = self.relu(self.bn1(self.conv1(point_cloud))) out2 = self.relu(self.bn2(self.conv2(out1))) out3 = self.relu(self.bn3(self.conv3(out2))) trans_feat = self.fstn(out3) x = out3.transpose(0, 2, 1) net_transformed = nn.bmm(x, trans_feat) net_transformed = net_transformed.transpose(0, 2, 1) out4 = self.relu(self.bn4(self.conv4(net_transformed))) out5 = self.bn5(self.conv5(out4)) out_max = jt.argmax(out5, 2, keepdims=True)[1] out_max = out_max.view(-1, 2048) out_max = concat((out_max, label), 1) expand = out_max.view(-1, 2048 + 16, 1).repeat(1, 1, N) concat_feature = concat([expand, out1, out2, out3, out4, out5], 1) net = self.relu(self.bns1(self.convs1(concat_feature))) net = self.relu(self.bns2(self.convs2(net))) net = self.relu(self.bns3(self.convs3(net))) net = self.convs4(net) return net
def test(): model.eval() logits, accs = model(), [] for _, mask in data('train_mask', 'val_mask', 'test_mask'): y_ = data.y[mask] tmp = [] for i in range(mask.shape[0]): if mask[i] == True: tmp.append(logits[i]) logits_ = jt.stack(tmp) pred, _ = jt.argmax(logits_, dim=1) acc = pred.equal(y_).sum().item() / mask.sum().item() accs.append(acc) return accs
def align_and_update_state_dicts(model_state_dict, loaded_state_dict): """ Strategy: suppose that the models that we will create will have prefixes appended to each of its keys, for example due to an extra level of nesting that the original pre-trained weights from ImageNet won't contain. For example, model.state_dict() might return backbone[0].body.res2.conv1.weight, while the pre-trained model contains res2.conv1.weight. We thus want to match both parameters together. For that, we look for each model weight, look among all loaded keys if there is one that is a suffix of the current weight name, and use it if that's the case. If multiple matches exist, take the one with longest size of the corresponding name. For example, for the same model as before, the pretrained weight file can contain both res2.conv1.weight, as well as conv1.weight. In this case, we want to match backbone[0].body.conv1.weight to conv1.weight, and backbone[0].body.res2.conv1.weight to res2.conv1.weight. """ current_keys = sorted(list(model_state_dict.keys())) loaded_keys = sorted(list(loaded_state_dict.keys())) # get a matrix of string matches, where each (i, j) entry correspond to the size of the # loaded_key string, if it matches match_matrix = [ len(j) if i.endswith(j) else 0 for i in current_keys for j in loaded_keys ] match_matrix = jt.array(match_matrix).reshape(len(current_keys), len(loaded_keys)) idxs, max_match_size = jt.argmax(match_matrix, 1) # remove indices that correspond to no-match idxs[max_match_size == 0] = -1 # used for logging max_size = max([len(key) for key in current_keys]) if current_keys else 1 max_size_loaded = max([len(key) for key in loaded_keys]) if loaded_keys else 1 log_str_template = "{: <{}} loaded from {: <{}} of shape {}" logger = logging.getLogger(__name__) for idx_new, idx_old in enumerate(idxs.numpy().tolist()): if idx_old == -1: continue key = current_keys[idx_new] key_old = loaded_keys[idx_old] model_state_dict[key] = loaded_state_dict[key_old] logger.info( log_str_template.format( key, max_size, key_old, max_size_loaded, tuple(loaded_state_dict[key_old].shape), ))
def farthest_point_sample(xyz, npoint): """ Input: xyz: pointcloud data, [B, N, C] npoint: number of samples Return: centroids: sampled pointcloud index, [B, npoint] """ #import ipdb; ipdb.set_trace() #device = xyz.device B, N, C = xyz.shape centroids = jt.zeros((B, npoint)) distance = jt.ones((B, N)) * 1e10 farthest = np.random.randint(0, N, B, dtype='l') batch_indices = np.arange(B, dtype='l') farthest = jt.array(farthest) batch_indices = jt.array(batch_indices) # jt.sync_all(True) # print (xyz.shape, farthest.shape, batch_indices.shape, centroids.shape, distance.shape) for i in range(npoint): centroids[:, i] = farthest centroid = xyz[batch_indices, farthest, :] centroid = centroid.view(B, 1, 3) dist = jt.sum((xyz - centroid.repeat(1, N, 1))**2, 2) mask = dist < distance # distance = mask.ternary(distance, dist) # print (mask.size()) if mask.sum().data[0] > 0: distance[mask] = dist[mask] # bug if mask.sum() == 0 farthest = jt.argmax(distance, 1)[0] # print (farthest) # print (farthest.shape) # B, N, C = xyz.size() # sample_list = random.sample(range(0, N), npoint) # centroids = jt.zeros((1, npoint)) # centroids[0,:] = jt.array(sample_list) # centroids = centroids.view(1, -1).repeat(B, 1) # x_center = x[:,sample_list, :] return centroids
def test(net, test_data, state): net.eval() loss_avg = 0.0 correct = 0 start_time = time.time() for batch_idx, (data, target) in enumerate(test_data): data, target = jt.array(data), jt.array(target) # forward output = net(data) loss = jt.nn.cross_entropy_loss(output, target) # accuracy pred = jt.argmax(output, dim=1)[0] correct += float(jt.sum(pred == target).data[0]) # test loss average loss_avg += float(loss.data[0]) end_time = time.time() fps = (len(test_data) * test_data.batch_size) / (end_time - start_time) state['test_loss'] = loss_avg / len(test_data) state['test_accuracy'] = correct / (len(test_data) * test_data.batch_size) state['test_fps'] = fps
def log_prob(self, x): x = jt.argmax(x, dim=-1)[0] return Categorical.log_prob(self, x)
def train(model): batch_size = 16 train_loader = ShapeNetPart(partition='trainval', num_points=2048, class_choice=None, batch_size=batch_size, shuffle=True) test_loader = ShapeNetPart(partition='test', num_points=2048, class_choice=None, batch_size=batch_size, shuffle=False) seg_num_all = 50 seg_start_index = 0 print(str(model)) base_lr = 0.01 optimizer = nn.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=1e-4) lr_scheduler = LRScheduler(optimizer, base_lr) # criterion = nn.cross_entropy_loss() # here best_test_iou = 0 for epoch in range(200): #################### # Train #################### lr_scheduler.step(len(train_loader) * batch_size) train_loss = 0.0 count = 0.0 model.train() train_true_cls = [] train_pred_cls = [] train_true_seg = [] train_pred_seg = [] train_label_seg = [] # debug = 0 for data, label, seg in train_loader: # with jt.profile_scope() as report: seg = seg - seg_start_index label_one_hot = np.zeros((label.shape[0], 16)) # print (label.size()) for idx in range(label.shape[0]): label_one_hot[idx, label.numpy()[idx, 0]] = 1 label_one_hot = jt.array(label_one_hot.astype(np.float32)) data = data.permute(0, 2, 1) # for pointnet it should not be committed batch_size = data.size()[0] # print ('input data shape') # print (data.shape, label_one_hot.shape) # for pointnet b c n for pointnet2 b n c seg_pred = model(data, label_one_hot) seg_pred = seg_pred.permute(0, 2, 1) # print (seg_pred.size()) # print (seg_pred.size(), seg.size()) loss = nn.cross_entropy_loss(seg_pred.view(-1, seg_num_all), seg.view(-1)) # print (loss.data) optimizer.step(loss) pred = jt.argmax(seg_pred, dim=2)[0] # (batch_size, num_points) # print ('pred size =', pred.size(), seg.size()) count += batch_size train_loss += loss.numpy() * batch_size seg_np = seg.numpy() # (batch_size, num_points) pred_np = pred.numpy() # (batch_size, num_points) # print (type(label)) label = label.numpy() # added train_true_cls.append( seg_np.reshape(-1)) # (batch_size * num_points) train_pred_cls.append( pred_np.reshape(-1)) # (batch_size * num_points) train_true_seg.append(seg_np) train_pred_seg.append(pred_np) temp_label = label.reshape(-1, 1) train_label_seg.append(temp_label) # print(report) train_true_cls = np.concatenate(train_true_cls) train_pred_cls = np.concatenate(train_pred_cls) # print (train_true_cls.shape ,train_pred_cls.shape) train_acc = metrics.accuracy_score(train_true_cls, train_pred_cls) avg_per_class_acc = metrics.balanced_accuracy_score( train_true_cls.data, train_pred_cls.data) # print ('train acc =',train_acc, 'avg_per_class_acc', avg_per_class_acc) train_true_seg = np.concatenate(train_true_seg, axis=0) # print (len(train_pred_seg), train_pred_seg[0].shape) train_pred_seg = np.concatenate(train_pred_seg, axis=0) # print (len(train_label_seg), train_label_seg[0].size()) # print (train_label_seg[0]) train_label_seg = np.concatenate(train_label_seg, axis=0) # print (train_pred_seg.shape, train_true_seg.shape, train_label_seg.shape) train_ious = calculate_shape_IoU(train_pred_seg, train_true_seg, train_label_seg, None) outstr = 'Train %d, loss: %.6f, train acc: %.6f, train avg acc: %.6f, train iou: %.6f' % ( epoch, train_loss * 1.0 / count, train_acc, avg_per_class_acc, np.mean(train_ious)) # io.cprint(outstr) print(outstr) #################### # Test #################### test_loss = 0.0 count = 0.0 model.eval() test_true_cls = [] test_pred_cls = [] test_true_seg = [] test_pred_seg = [] test_label_seg = [] for data, label, seg in test_loader: seg = seg - seg_start_index label_one_hot = np.zeros((label.shape[0], 16)) for idx in range(label.shape[0]): label_one_hot[idx, label.numpy()[idx, 0]] = 1 label_one_hot = jt.array(label_one_hot.astype(np.float32)) data = data.permute(0, 2, 1) # for pointnet should not be commit batch_size = data.size()[0] seg_pred = model(data, label_one_hot) seg_pred = seg_pred.permute(0, 2, 1) loss = nn.cross_entropy_loss(seg_pred.view(-1, seg_num_all), seg.view(-1, 1).squeeze(-1)) pred = jt.argmax(seg_pred, dim=2)[0] count += batch_size test_loss += loss.numpy() * batch_size seg_np = seg.numpy() pred_np = pred.numpy() label = label.numpy() # added test_true_cls.append(seg_np.reshape(-1)) test_pred_cls.append(pred_np.reshape(-1)) test_true_seg.append(seg_np) test_pred_seg.append(pred_np) test_label_seg.append(label.reshape(-1, 1)) test_true_cls = np.concatenate(test_true_cls) test_pred_cls = np.concatenate(test_pred_cls) test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls) avg_per_class_acc = metrics.balanced_accuracy_score( test_true_cls, test_pred_cls) test_true_seg = np.concatenate(test_true_seg, axis=0) test_pred_seg = np.concatenate(test_pred_seg, axis=0) test_label_seg = np.concatenate(test_label_seg) test_ious = calculate_shape_IoU(test_pred_seg, test_true_seg, test_label_seg, None) outstr = 'Test %d, loss: %.6f, test acc: %.6f, test avg acc: %.6f, test iou: %.6f' % ( epoch, test_loss * 1.0 / count, test_acc, avg_per_class_acc, np.mean(test_ious)) print(outstr)
def sample_images(batches_done): """Saves a generated sample from the validation set""" batches = next(iter(val_dataloader)) real_A = batches[0] real_A_eyel = batches[1] real_A_eyer = batches[2] real_A_nose = batches[3] real_A_mouth = batches[4] real_A_hair = batches[5] real_A_bg = batches[6] mask = batches[7] mask2 = batches[8] center = batches[9] cmaskel = batches[10] cmasker = batches[11] cmaskno = batches[12] cmaskmo = batches[13] maskh = mask * mask2 maskb = inverse_mask(mask2) fake_B0 = G_global(real_A) # EYES, NOSE, MOUTH fake_B_eyel1 = G_l_eyel(real_A_eyel) fake_B_eyel2 = ae_eyel(fake_B_eyel1) fake_B_eyel = add_with_mask(fake_B_eyel2, fake_B_eyel1, cmaskel) fake_B_eyer1 = G_l_eyer(real_A_eyer) fake_B_eyer2 = ae_eyer(fake_B_eyer1) fake_B_eyer = add_with_mask(fake_B_eyer2, fake_B_eyer1, cmasker) fake_B_nose1 = G_l_nose(real_A_nose) fake_B_nose2 = ae_nose(fake_B_nose1) fake_B_nose = add_with_mask(fake_B_nose2, fake_B_nose1, cmaskno) fake_B_mouth1 = G_l_mouth(real_A_mouth) outputs1 = CLm(real_A_mouth) pred = jt.argmax(outputs1, dim=1)[0] fake_B_mouth2w = ae_mowhite(fake_B_mouth1) fake_B_mouth2b = ae_moblack(fake_B_mouth1) fake_B_mouth2s = jt.contrib.concat((fake_B_mouth2w, fake_B_mouth2b), 1) bs, c, h, w = fake_B_mouth2s.shape index = pred + jt.arange(bs) * c fake_B_mouth2 = fake_B_mouth2s.reshape([-1, h, w])[index].reshape([bs, 1, h, w]) fake_B_mouth = add_with_mask(fake_B_mouth2, fake_B_mouth1, cmaskmo) # HAIR & BG outputs2 = CLh(real_A_hair) onehot2 = getonehot(outputs2, 3, bs) fake_B_hair = G_l_hair(real_A_hair, onehot2) fake_B_bg = G_l_bg(real_A_bg) # PARTCOMBINE fake_B1 = partCombiner2_bg(center, fake_B_eyel, fake_B_eyer, fake_B_nose, fake_B_mouth, fake_B_hair, fake_B_bg, maskh, maskb, comb_op=1, load_h=opt.img_height, load_w=opt.img_width) # FUSION NET fake_B = G_combine(jt.contrib.concat((fake_B0, fake_B1), 1)) img_sample = np.concatenate( [real_A.data, fake_B.repeat(1, 3, 1, 1).data], -2) save_image(img_sample, "images/%s/%s.jpg" % (opt.dataset_name, batches_done), nrow=5)
fake = jt.zeros([real_A.shape[0], 1]).stop_grad() fake_B0 = G_global(real_A) # EYES, NOSE, MOUTH fake_B_eyel1 = G_l_eyel(real_A_eyel) fake_B_eyel2 = ae_eyel(fake_B_eyel1) fake_B_eyel = add_with_mask(fake_B_eyel2, fake_B_eyel1, cmaskel) fake_B_eyer1 = G_l_eyer(real_A_eyer) fake_B_eyer2 = ae_eyer(fake_B_eyer1) fake_B_eyer = add_with_mask(fake_B_eyer2, fake_B_eyer1, cmasker) fake_B_nose1 = G_l_nose(real_A_nose) fake_B_nose2 = ae_nose(fake_B_nose1) fake_B_nose = add_with_mask(fake_B_nose2, fake_B_nose1, cmaskno) fake_B_mouth1 = G_l_mouth(real_A_mouth) outputs1 = CLm(real_A_mouth) pred = jt.argmax(outputs1, dim=1)[0] fake_B_mouth2w = ae_mowhite(fake_B_mouth1) fake_B_mouth2b = ae_moblack(fake_B_mouth1) fake_B_mouth2s = jt.contrib.concat((fake_B_mouth2w, fake_B_mouth2b), 1) bs, c, h, w = fake_B_mouth2s.shape index = pred + jt.arange(bs) * c fake_B_mouth2 = fake_B_mouth2s.reshape([-1, h, w])[index].reshape( [bs, 1, h, w]) fake_B_mouth = add_with_mask(fake_B_mouth2, fake_B_mouth1, cmaskmo) # HAIR & BG outputs2 = CLh(real_A_hair) onehot2 = getonehot(outputs2, 3, bs) fake_B_hair = G_l_hair(real_A_hair, onehot2) fake_B_bg = G_l_bg(real_A_bg) # PARTCOMBINE fake_B1 = partCombiner2_bg(center,
def getonehot(outputs, classes, batch_size): index = jt.argmax(outputs,1) y = jt.unsqueeze(index[0],1) onehot = jt.zeros([batch_size, classes]) onehot.scatter_(1, y, jt.array(1.)) return onehot
def log_prob(self, x): if len(x.shape) == 1: x = x.unsqueeze(0) logits = self.logits.broadcast(x.shape) indices = jt.argmax(x, dim=-1)[0] return logits.gather(1, indices.unsqueeze(-1)).reshape(-1)