点云生成：基于Paddle2.0实现WGAN-GP在点云上的一些尝试

!unzip data/data50045/modelnet40_normal_resampled.zip!mv modelnet40_normal_resampled dataset

import osimport numpy as npimport randomfrom mpl_toolkits import mplot3dimport matplotlib.pyplot as pltimport paddleimport paddle.nn.functional as Ffrom paddle.nn import Conv2D, Conv2DTranspose, MaxPool2D, Linear, BatchNorm, Dropout, ReLU, Tanh, LeakyReLU, Sequential

category = {    'airplane': 0,
}

def getDatalist(file_path='./dataset/modelnet40_shape_names.txt'):
    f = open(file_path, 'r')
    f_train = open('./dataset/train.txt', 'w')
    f_test = open('./dataset/test.txt', 'w')    for category in f:        if category.split('\n')[0] == 'airplane':
            dict_path = os.path.join('./dataset/', category.split('\n')[0])
            data_dict = os.listdir(dict_path)
            count = 0
            for data_path in data_dict:                if count % 61 != 0:
                    f_train.write(os.path.join(dict_path, data_path) + ' ' + category)                else:
                    f_test.write(os.path.join(dict_path, data_path) + ' ' + category)
                count += 1
    f_train.close()
    f_test.close()
    f.close()if __name__ == '__main__':
    getDatalist()

def pointDataLoader(file_path='./dataset/train.txt', mode='train'):
    BATCHSIZE = 8
    MAX_POINT = 250
    datas = []
    labels = []
    f = open(file_path)    for data_list in f:
        point_data = []
        data_path = data_list.split(' ')[0]
        data_file = open(data_path)
        point_num = 0
        for points in data_file:            if point_num == MAX_POINT:                break
            point_data.append([                float(points.split(',')[0]),                float(points.split(',')[1]),                float(points.split(',')[2])
            ])
            point_num += 1
        datas.append(point_data)
        labels.append(category[data_list.split(' ')[1].split('\n')[0]])
    f.close()
    datas = np.array(datas)
    labels = np.array(labels)

    index_list = list(range(len(datas)))    def pointDataGenerator():
        if mode == 'train':
            random.shuffle(index_list)
        datas_list = []
        labels_list = []        for i in index_list:
            data = np.reshape(datas[i], [1, 250, 3]).astype('float32')
            label = np.reshape(labels[i], [1]).astype('int64')
            datas_list.append(data) 
            labels_list.append(label)            if len(datas_list) == BATCHSIZE:                yield np.array(datas_list), np.array(labels_list)
                datas_list = []
                labels_list = []        if len(datas_list) > 0:            yield np.array(datas_list), np.array(labels_list)    return pointDataGenerator

class FeatureNet(paddle.nn.Layer):
    def __init__(self, name_scope='FeatureNet_', num_point=256):
        super(FeatureNet, self).__init__()
        self.input_transform_net = Sequential(
            Conv2D(1, 64, (1, 3)),
            BatchNorm(64),
            ReLU(),
            Conv2D(64, 256, (1, 1)),
            BatchNorm(256),
            ReLU(),
            MaxPool2D((num_point, 1))
        )
        self.input_fc = Sequential(
            Linear(256, 64),
            ReLU(),
            Linear(64, 9, 
                weight_attr=paddle.framework.ParamAttr(initializer=paddle.nn.initializer.Assign(paddle.zeros((64, 9)))),
                bias_attr=paddle.framework.ParamAttr(initializer=paddle.nn.initializer.Assign(paddle.reshape(paddle.eye(3), [-1])))
            )
        )
        self.mlp_1 = Sequential(
            Conv2D(1, 64, (1, 3)),
            BatchNorm(64),
            ReLU(),
            Conv2D(64, 16,(1, 1)),
            BatchNorm(16),
            ReLU(),
        )
        self.feature_transform_net = Sequential(
            Conv2D(16, 16, (1, 1)),
            BatchNorm(16),
            ReLU(),
            MaxPool2D((num_point, 1))
        )
        self.feature_fc = Sequential(
            Linear(16, 8),
            ReLU(),
            Linear(8, 16*16)
        )
        self.mlp_2 = Sequential(
            Conv2D(16, 8, (1, 1)),
            BatchNorm(8),
            ReLU()
        )    
    def forward(self, inputs):
        """
        input: [batchsize, 1, 250, 3]
        output: [batchsize, 250, 1]
        """
        batchsize = inputs.shape[0]

        t_net = self.input_transform_net(inputs)
        t_net = paddle.squeeze(t_net)
        t_net = self.input_fc(t_net)
        t_net = paddle.reshape(t_net, [batchsize, 3, 3])

        x = paddle.squeeze(inputs)
        x = paddle.matmul(x, t_net)
        x = paddle.unsqueeze(x, axis=1)
        x = self.mlp_1(x)

        t_net = self.feature_transform_net(x)
        t_net = paddle.squeeze(t_net)
        t_net = self.feature_fc(t_net)
        t_net = paddle.reshape(t_net, [batchsize, 16, 16])

        x = paddle.squeeze(x)
        x = paddle.transpose(x, (0, 2, 1))
        x = paddle.matmul(x, t_net)
        x = paddle.transpose(x, (0, 2, 1))
        x = paddle.unsqueeze(x, axis=-1)
        x = self.mlp_2(x)
        x = paddle.max(x, axis=2)        return x

class UFeatureNet(paddle.nn.Layer):
    def __init__(self, name_scope='UFeatureNet_', num_point=1024):
        super(UFeatureNet, self).__init__()
        self.stage_1 = Sequential(
            Conv2DTranspose(1, 4, (1, 3)),
            BatchNorm(4),
            LeakyReLU(),
            Conv2D(4, 16, (1, 1)),
            BatchNorm(16),
            LeakyReLU()
        )
        self.stage_2 = Sequential(
            Conv2DTranspose(16, 32, (4, 4), (2, 1)),
            BatchNorm(32),
            LeakyReLU(),
            Conv2D(32, 128, (1, 1)),
            BatchNorm(128),
            LeakyReLU()
        )
        self.stage_3 = Sequential(
            Conv2DTranspose(128, 128, (4, 4), (2, 1)),
            BatchNorm(128),
            LeakyReLU(),
            Conv2D(128, 64, (1, 3)),
            BatchNorm(64),
            LeakyReLU()
        )
        self.stage_4 = Sequential(
            Conv2DTranspose(64, 32, (4, 1), (1, 1)),
            BatchNorm(32),
            LeakyReLU(),
            Conv2D(32, 32, (1, 3)),
            BatchNorm(32),
            LeakyReLU()
        )
        self.stage_5 = Sequential(
            Conv2DTranspose(32, 16, (2, 1), (1, 1)),
            BatchNorm(16),
            LeakyReLU(),
            Conv2D(16, 1, (1, 3)),
            Tanh()
        )    def forward(self, inputs):
        """
        input: [batchsize, 1, 60, 1]
        output: [batchsize, 1, 250, 3]
        """
        x = self.stage_1(inputs)
        x = self.stage_2(x)
        x = self.stage_3(x)
        x = self.stage_4(x)
        x = self.stage_5(x)        return x

class D(paddle.nn.Layer):
    def __init__(self, name_scope='D_'):
        super(D, self).__init__()
        self.feature_net = FeatureNet()
        self.fc = Sequential(
            Linear(8, 8),
            ReLU(),
            Dropout(p=0.7),
            Linear(8, 1)
        )    def forward(self, inputs):
        """
        input: [batchsize, 1, 250, 3]
        output: [batchsize, 1]
        """
        x = self.feature_net(inputs)
        x = paddle.squeeze(x)
        x = self.fc(x)        return x

class G(paddle.nn.Layer):
    def __init__(self, name_scope='G_'):
        super(G, self).__init__()
        self.u_feature_net = UFeatureNet()    def forward(self, inputs):
        """
        input: [batchsize, 1, 60, 1]
        output: [batchsize, 1, 250, 3]
        """
        x = self.u_feature_net(inputs)        return x

Discriminator = D()
paddle.summary(Discriminator, (8, 1, 250, 3))

Generator = G()
paddle.summary(Generator, (8, 1, 60, 1))

def gradient_penalty(discriminator, real, fake, batchsize):
    t = paddle.uniform((batchsize,1,1,1))
    t = paddle.expand_as(t, real)
    inter = t * real +  (1-t) * fake
    inter.stop_gradient = False
    inter_ = discriminator(inter)
    grads = paddle.grad(inter_, [inter])[0]
    grads = paddle.reshape(grads, [batchsize, grads.shape[1], grads.shape[2], grads.shape[3]])
    epsilon = 1e-12
    norm = paddle.sqrt(
        paddle.mean(paddle.square(grads), axis=1) + epsilon
    )
    gp = paddle.mean((norm - 1)**2) * 10

    return gp

def draw(data):
    zdata = []
    xdata = []
    ydata = []    for i in data[0][0]:
        xdata.append(i[0])
        ydata.append(i[1])
        zdata.append(i[2])
    xdata = np.array(xdata)
    ydata = np.array(ydata)
    zdata = np.array(zdata)

    ax = plt.axes(projection='3d')
    ax.scatter3D(xdata, ydata, zdata, c='r')
    plt.savefig('fake.png')

def train():
    train_loader = pointDataLoader(file_path='./dataset/train.txt', mode='train')
    Discriminator = D()
    Generator = G()

    Discriminator.train()
    Generator.train()

    optim1 = paddle.optimizer.Adam(parameters=Discriminator.parameters(), weight_decay=0.001, learning_rate=1e-5)
    optim2 = paddle.optimizer.Adam(parameters=Generator.parameters(), weight_decay=0.001, learning_rate=1e-5)

    epoch_num = 2000
    train_g = 2
    for epoch in range(epoch_num):        for batch_id, data in enumerate(train_loader()):
            real = paddle.to_tensor(data[0])
            noise = paddle.uniform((real.shape[0], 1, 60, 1))
            fake = Generator(noise)

            fake_loss = paddle.mean(Discriminator(fake))
            real_loss = -1*paddle.mean(Discriminator(real))
            gp = gradient_penalty(Discriminator, real, fake, real.shape[0])
            d_loss = fake_loss + real_loss + gp

            d_loss.backward()
            optim1.step()
            optim1.clear_grad()            for _ in range(train_g):
                noise = paddle.uniform((real.shape[0], 1, 60, 1))
                fake = Generator(noise)
                g_loss = -1*paddle.mean(Discriminator(fake))

                g_loss.backward()
                optim2.step()
                optim2.clear_grad()            if batch_id % 32 == 0: 
                print("epoch: {}, batch_id: {}, d_loss is: {}, g_loss is: {}".format(epoch, batch_id, d_loss.numpy(), g_loss.numpy()))        
        if epoch % 2 == 0:
            noise = paddle.uniform((real.shape[0], 1, 60, 1))
            fake = Generator(noise)
            draw(fake.numpy())        if epoch % 20 == 0:
            paddle.save(Discriminator.state_dict(), './model/D.pdparams')
            paddle.save(optim1.state_dict(), './model/D.pdopt')
            paddle.save(Generator.state_dict(), './model/G.pdparams')
            paddle.save(optim2.state_dict(), './model/G.pdopt')if __name__ == '__main__':
    train()