NolenChen / 3DStructurePoints

77 stars 15 forks source link

简化了模型代码 #5

Open QiangZiBro opened 4 years ago

QiangZiBro commented 4 years ago

1.使用Conv1dProbLayer简化了模型代码

class Conv1dProbLayer(BaseModel):
    def __init__(self, in_channels, out_channels, out=False, kernel_size=1, dropout=0.2):
        super().__init__()
        self.out = out
        self.dropout_conv_bn_layer = nn.Sequential(
            nn.Dropout(dropout),
            nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size),
            nn.BatchNorm1d(num_features=out_channels),
        )
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=2)

    def forward(self, x):
        x = self.dropout_conv_bn_layer(x)
        if self.out:
            x = self.softmax(x)
        else:
            x = self.relu(x)
        return x
  1. readme 使用指定cuda版本安装torch