PointNet 分类部分复现
只用了modelnet40的第一个部分训练的
大致思路和图像一样,只不过用的conv1d,在数据集处理上也稍有不同
my_Dataset.py:
![](https://images.cnblogs.com/OutliningIndicators/ContractedBlock.gif)
import numpy as np import h5py import torch import random import torchvision.transforms as transforms from torch.utils import data #dataset是否要做个转置呢并变成tensor class Dataset(data.Dataset): def __init__(self, root): super(Dataset, self).__init__() file = h5py.File(root, 'r') self.data = file['data'][:] #label要展开并变成一维的 self.label = file['label'][:].reshape(-1, file['label'].shape[0]).squeeze(0) def __getitem__(self, index): #一个点云坐标输入进来时是n * 3,不像图片是n * m * 3,点云坐标没有高度,x、y、z分别代表三个通道 #图片用ToTensor变成3 * n * m, 点云坐标转置一下就行了,变成3 * n,然后再变成tensor return torch.tensor(self.data[index].T), self.label[index] def __len__(self): return len(self.label)
configuratiion.py:
![](https://images.cnblogs.com/OutliningIndicators/ContractedBlock.gif)
import torch class config(): batch_size = 4 num_epochs = 10 num_classes = 40 num_workers = 8 checkpoints_root = 'C:/Users/Dell/PycharmProjects/PointNet/checkpoints' device = 'cuda' if torch.cuda.is_available() else 'cpu' log_dir = 'C:/Users/Dell/PycharmProjects/PointNet/checkpoints/log'
Model.py:
![](https://images.cnblogs.com/OutliningIndicators/ContractedBlock.gif)
import torch import torch.nn as nn import torch.nn.functional as F import numpy as np from tqdm import tqdm from configuration import config device = 'cuda' if torch.cuda.is_available() else 'cpu' con = config() #T-Net:将点云传进来,生成一个矩阵,+E后return class T_Net(nn.Module): def __init__(self, k): super().__init__() self.k = k self.conv1 = nn.Conv1d(self.k, 64, 1) self.conv2 = nn.Conv1d(64, 128, 1) self.conv3 = nn.Conv1d(128, 1024, 1) self.fc1 = nn.Linear(1024, 512) self.fc2 = nn.Linear(512, 256) self.fc3 = nn.Linear(256, self.k * self.k) self.bn1 = nn.BatchNorm1d(64) self.bn2 = nn.BatchNorm1d(128) self.bn3 = nn.BatchNorm1d(1024) self.bn4 = nn.BatchNorm1d(512) self.bn5 = nn.BatchNorm1d(256) def forward(self, x): bs = x.size(0) x = F.relu(self.bn1(self.conv1(x))) x = F.relu(self.bn2(self.conv2(x))) x = F.relu(self.bn3(self.conv3(x))) x = F.max_pool1d(x, x.size(-1)) x = x.view(x.size(0), -1) x = F.relu(self.bn4(self.fc1(x))) x = F.relu(self.bn5(self.fc2(x))) x = self.fc3(x) #生成一个3 * 3单位矩阵E,并将其扩充为bs个3 * 3的,repeat中的两个1表示3 * 1,3 * 1, #即为两个系数k、m,使得扩充后行列分别为,行 * k, 列 * m E = torch.eye(self.k, requires_grad = True).repeat(bs, 1, 1) E.to(device) matrix = x.view(-1, self.k, self.k) + E return matrix class PNet(nn.Module): def __init__(self): super(PNet, self).__init__() self.input_transform = T_Net(k = 3) self.feature_transform = T_Net(k = 64) self.conv1 = nn.Conv1d(3, 64, 1) self.conv2 = nn.Conv1d(64, 128, 1) self.conv3 = nn.Conv1d(128, 1024, 1) self.fc1 = nn.Linear(1024, 512) self.fc2 = nn.Linear(512, 256) self.fc3 = nn.Linear(256, con.num_classes) self.bn1 = nn.BatchNorm1d(64) self.bn2 = nn.BatchNorm1d(128) self.bn3 = nn.BatchNorm1d(1024) self.bn4 = nn.BatchNorm1d(512) self.bn5 = nn.BatchNorm1d(256) self.dropout = nn.Dropout(0.3) def forward(self, x): input_matrix = self.input_transform(x) x = torch.bmm(torch.transpose(x, 1, 2), input_matrix).transpose(1, 2) x = F.relu(self.bn1(self.conv1(x))) feature_matrix = self.feature_transform(x) x = torch.bmm(x.transpose(1, 2), feature_matrix).transpose(1, 2) x = F.relu(self.bn2(self.conv2(x))) x = self.bn3(self.conv3(x)) #3 * n的点云坐标经过一系列卷积层之后,变成了1024 * n的 #经过max_pool求1024维的每一维最大值,变成了1024 * 1 #max_pool1d的第二个参数表示池化的范围,当然是n,-1表示size的倒数第一个 x = F.max_pool1d(x, x.size(-1)) x = x.view(x.size(0), -1) x = F.relu(self.bn4(self.fc1(x))) x = F.relu(self.bn5(self.dropout(self.fc2(x)))) x = self.fc3(x) return x, input_matrix, feature_matrix
test.py:
![](https://images.cnblogs.com/OutliningIndicators/ContractedBlock.gif)
import torch import torch.nn import torch.utils.data.dataloader as Dataloader from configuration import config from my_Dataset import Dataset from Model import PNet import os if __name__ == '__main__': con = config() model = PNet() checkpoint = torch.load(os.path.join(con.checkpoints_root, 'checkpoint_10.pkl')) model.load_state_dict(checkpoint['model']) model.to(con.device) dataset = Dataset('H:/DataSet/modelnet40_ply_hdf5_2048/ply_data_test0.h5') dataloader = Dataloader.DataLoader(dataset, batch_size=2, shuffle = True) cnt = 0 for data, label in dataloader: data.to(con.device) output = model(data)[0] pred = torch.max(output, 1)[1] print(pred, label) cnt += 1 if cnt == 20: break
train.py:
![](https://images.cnblogs.com/OutliningIndicators/ContractedBlock.gif)
import torch import torch.nn as nn import torch.utils.data.dataloader as Dataloader from tqdm import tqdm from my_Dataset import Dataset from Model import PNet from configuration import config import os from tensorboardX import SummaryWriter con = config() def loss_funtion(output, label, input_matrix, featrue_matrix, alpha = 0.0001): loss = nn.CrossEntropyLoss() bs = output.size(0) E_3 = torch.eye(3, requires_grad = True).repeat(bs, 1, 1) E_64 = torch.eye(64, requires_grad = True).repeat(bs, 1, 1) E_3.to(con.device) E_64.to(con.device) diff3 = E_3 - torch.bmm(input_matrix, input_matrix.transpose(1, 2)) diff64 = E_64 - torch.bmm(featrue_matrix, featrue_matrix.transpose(1, 2)) #注意label要是int64类型 label = label.type(torch.LongTensor) return loss(output, label) + alpha * (torch.norm(diff3) + torch.norm(diff64)) / float(bs) if __name__ == '__main__': data_path = 'H:/DataSet/modelnet40_ply_hdf5_2048/ply_data_train0.h5' dataset = Dataset(data_path) dataloader = Dataloader.DataLoader(dataset, batch_size = con.batch_size, shuffle = True, num_workers = con.num_workers) model = PNet() model.to(con.device) optimizer = torch.optim.Adam(model.parameters(), lr = 0.01) tbwriter = SummaryWriter(logdir = con.log_dir) for epoch in range(con.num_epochs): total_loss = 0 total_true = 0 cnt = 0 total_img = 0 for data, label in tqdm(dataloader): data.to(con.device) label.to(con.device) optimizer.zero_grad() output, input_matrix, feature_matrix = model(data) loss_value = loss_funtion(output, label, input_matrix, feature_matrix) loss_value.backward() optimizer.step() print(type(output), output) pred = torch.max(output, 1)[1] total_true += torch.sum(pred == label) total_loss += loss_value cnt += 1 total_img += len(label) tbwriter.add_scalar('Loss', total_loss / float(cnt), epoch) tbwriter.add_scalar('Accuracy', total_true / float(total_img), epoch) print('Loss:{:.4f}, Accuracy:{:.2f}'.format(total_loss / float(cnt), total_true / float(total_img))) if (epoch + 1) % 10 == 0: state = { 'model': model.state_dict() } torch.save(state, os.path.join(con.checkpoints_root, 'checkpoint_{}.pkl'.format(epoch + 1))) print('Train Accepted')
自己选择的路,跪着也要走完。朋友们,虽然这个世界日益浮躁起来,只要能够为了当时纯粹的梦想和感动坚持努力下去,不管其它人怎么样,我们也能够保持自己的本色走下去。