GAN模型迁移适配指导

1. 模型概述

GAN生成对抗网络是一种生成模型，它由两个部分组成，分别称为“生成器”和“判别器”。生成器将随机噪声作为输入，并将其转换为与训练数据相似的输出。判别器则接收一组样本作为输入，并试图区分真实的训练样本与生成器创建的样本。两者会同时进行训练：判别器努力越来越擅长辨别真假数据，而生成器则努力越来越擅长欺骗判别器。该模型在药物发现、材料科学和计算化学等任务中展现出强大能力，本文描述的GAN模型是基于DeepChem套件实现的，后续适配也是基于该套件修改。

2. 准备运行环境

2.1 软件环境

组件	版本
Python	3.10.19
PyTorch	2.1.0
torch_npu	2.1.0.post13
CANN	8.1.RC1

2.2 硬件环境

设备型号	NPU 配置
Atlas 800T A2	单卡 / 多卡

2.3 准备镜像

镜像环境	镜像地址
公网	swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_2_1_ascend:pytorch_2.1.0-cann_8.1.rc1-py_3.10-euler_2.10.11-aarch64-snt9b-20250603154214-4e60e43

2.4 启动镜像

docker run -u root --privileged \
 --name {container_name} \
 --device /dev/davinci0 \
 --device /dev/davinci_manager \
 --device /dev/devmm_svm \
 --device /dev/hisi_hdc \
 -v /usr/local/dcmi:/usr/local/dcmi \
 -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
 -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
 -v /etc/ascend_install.info:/etc/ascend_install.info \
 -itd {image_id} /bin/bash

3 运行指导

3.1 安装已完成迁移的deepchem-ascend

docker exec -it {container_name} bash

直接安装 deepchem-ascend 二进制包，该包已基于适配代码重新编译并上传至 PyPI。

pip install deepchem-ascend==0.0.1

3.2 安装依赖

安装运行必要的依赖库

pip install torch-geometric

3.3 测试验证

测试代码举例如下：

3.6 运行测试代码

import numpy as np
import tempfile
from typing import Callable, Any, List, Optional, Tuple, Union
from deepchem.models.torch_models.torch_model import is_npu_available

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from deepchem.models.torch_models import GANModel

    # helper classes that depend on torch, they need to be in the try/catch block
    class Generator(nn.Module):
        """A simple generator for testing."""

        def __init__(self, noise_input_shape, conditional_input_shape):
            super(Generator, self).__init__()
            self.noise_input_shape = noise_input_shape
            self.conditional_input_shape = conditional_input_shape

            self.noise_dim = noise_input_shape[1:]
            self.conditional_dim = conditional_input_shape[1:]

            input_dim = sum(self.noise_dim) + sum(self.conditional_dim)

            # 适配npu：input_dim=3 扩展到 4（4 的倍数）
            self.expand = nn.Linear(input_dim, 4)
            self.output = nn.Linear(4, 1)

            # 处理 device：支持 'npu'
            if torch.cuda.is_available():
                device = torch.device('cuda')
            elif torch.backends.mps.is_available():
                device = torch.device('mps')
            elif is_npu_available():  # 关键：支持 NPU
                device = torch.device('npu')
            else:
                device = torch.device('cpu')

            # 修复：强制移动到 NPU
            self.to(device)

        def forward(self, input):
            noise_input, conditional_input = input

            inputs = torch.cat((noise_input, conditional_input), dim=1)
            # 解决npu非连续内存问题
            inputs = inputs.contiguous()

            # 扩展维度
            expanded = self.expand(inputs)  # [B, 4]

            output = self.output(expanded)
            return output

    class Discriminator(nn.Module):
        """A simple discriminator for testing."""

        def __init__(self, data_input_shape, conditional_input_shape):
            super(Discriminator, self).__init__()
            self.data_input_shape = data_input_shape
            self.conditional_input_shape = conditional_input_shape

            data_dim = data_input_shape[
                1:]  # Extracting the actual data dimension
            conditional_dim = conditional_input_shape[
                1:]  # Extracting the actual conditional dimension
            input_dim = sum(data_dim) + sum(conditional_dim)

            # Define the dense layers
            self.dense1 = nn.Linear(input_dim, 10)
            self.dense2 = nn.Linear(10, 1)

        def forward(self, input):
            data_input, conditional_input = input
            discrim_in = torch.cat((data_input, conditional_input), dim=1)
            x = F.relu(self.dense1(discrim_in))
            a = self.dense2(x)
            output = torch.sigmoid(a)
            return output

    class ExampleGANModel(GANModel):
        """A simple GAN for testing."""

        def get_noise_input_shape(self):
            return (
                100,
                2,
            )

        def get_data_input_shapes(self):
            return [(
                100,
                1,
            )]

        def get_conditional_input_shapes(self):
            return [(
                100,
                1,
            )]

        def create_generator(self):
            noise_dim = self.get_noise_input_shape()
            conditional_dim = self.get_conditional_input_shapes()[0]

            return nn.Sequential(Generator(noise_dim, conditional_dim))

        def create_discriminator(self):
            data_input_shape = self.get_data_input_shapes()[0]
            conditional_input_shape = self.get_conditional_input_shapes()[0]

            return nn.Sequential(
                Discriminator(data_input_shape, conditional_input_shape))

        # 必须重写：支持 NPU
        def __init__(self,
                    n_generators: int = 1,
                    n_discriminators: int = 1,
                    create_discriminator_loss: Optional[Callable] = None,
                    create_generator_loss: Optional[Callable] = None,
                    _call_discriminator: Optional[Callable] = None,
                    device: Optional[Union[str, torch.device]] = None,
                    **kwargs):
        
            # 处理 device：支持 'npu'
            if device is None:
                if torch.cuda.is_available():
                    self.device = torch.device('cuda')
                elif torch.backends.mps.is_available():
                    self.device = torch.device('mps')
                elif is_npu_available():  # 关键：支持 NPU
                    self.device = torch.device('npu')
                else:
                    self.device = torch.device('cpu')
            else:
                if isinstance(device, str):
                    if device.lower() == 'npu':
                        if not torch.npu.is_available():
                            raise RuntimeError("NPU is not available. Please install CANN and torch.npu.")
                        self.device = torch.device('npu')
                    else:
                        self.device = torch.device(device)
                else:
                    self.device = device
            
            print(f"模型所在设备: {self.device}")

            # 保存参数
            self.n_generators = n_generators
            self.n_discriminators = n_discriminators

            # 调用父类构造函数（传入 device）
            super().__init__(
                n_generators=n_generators,
                n_discriminators=n_discriminators,
                create_discriminator_loss=create_discriminator_loss,
                create_generator_loss=create_generator_loss,
                _call_discriminator=_call_discriminator,
                device=self.device,
                **kwargs
            )

    has_torch = True
except ModuleNotFoundError:
    has_torch = False

def create_generator(noise_dim, conditional_dim):
    noise_dim = noise_dim
    conditional_dim = conditional_dim[0]
    return nn.Sequential(Generator(noise_dim, conditional_dim))

def create_discriminator(data_input_shape, conditional_input_shape):
    data_input_shape = data_input_shape[0]
    conditional_input_shape = conditional_input_shape[0]
    return nn.Sequential(
        Discriminator(data_input_shape, conditional_input_shape))

def generate_batch(batch_size):
    """Draw training data from a Gaussian distribution, where the mean  is a conditional input."""
    means = 10 * np.random.random([batch_size, 1])
    values = np.random.normal(means, scale=2.0)
    return means, values

def generate_data(gan, batches, batch_size):
    for _ in range(batches):
        means, values = generate_batch(batch_size)
        batch = {gan.data_inputs[0]: values, gan.conditional_inputs[0]: means}
        yield batch

gan = ExampleGANModel(learning_rate=0.01)
data = generate_data(gan, 500, 100)
gan.fit_gan(data, generator_steps=0.5, checkpoint_interval=0)

# See if it has done a plausible job of learning the distribution.

means = 10 * np.random.random([1000, 1])
values = gan.predict_gan_generator(conditional_inputs=[means])

测试结果：

1. 模型概述

组件

版本

Python

3.10.19

PyTorch

2.1.0

torch_npu

2.1.0.post13

CANN

8.1.RC1

设备型号

NPU 配置

Atlas 800T A2

单卡 / 多卡

镜像环境

镜像地址

公网

swr.cn-southwest-2.myhuaweicloud.com/atelier/pytorch_2_1_ascend:pytorch_2.1.0-cann_8.1.rc1-py_3.10-euler_2.10.11-aarch64-snt9b-20250603154214-4e60e43

2.4 启动镜像

docker run -u root --privileged \
 --name {container_name} \
 --device /dev/davinci0 \
 --device /dev/davinci_manager \
 --device /dev/devmm_svm \
 --device /dev/hisi_hdc \
 -v /usr/local/dcmi:/usr/local/dcmi \
 -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
 -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
 -v /etc/ascend_install.info:/etc/ascend_install.info \
 -itd {image_id} /bin/bash

3.6 运行测试代码

import numpy as np
import tempfile
from typing import Callable, Any, List, Optional, Tuple, Union
from deepchem.models.torch_models.torch_model import is_npu_available

try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from deepchem.models.torch_models import GANModel

    # helper classes that depend on torch, they need to be in the try/catch block
    class Generator(nn.Module):
        """A simple generator for testing."""

        def __init__(self, noise_input_shape, conditional_input_shape):
            super(Generator, self).__init__()
            self.noise_input_shape = noise_input_shape
            self.conditional_input_shape = conditional_input_shape

            self.noise_dim = noise_input_shape[1:]
            self.conditional_dim = conditional_input_shape[1:]

            input_dim = sum(self.noise_dim) + sum(self.conditional_dim)

            # 适配npu：input_dim=3 扩展到 4（4 的倍数）
            self.expand = nn.Linear(input_dim, 4)
            self.output = nn.Linear(4, 1)

            # 处理 device：支持 'npu'
            if torch.cuda.is_available():
                device = torch.device('cuda')
            elif torch.backends.mps.is_available():
                device = torch.device('mps')
            elif is_npu_available():  # 关键：支持 NPU
                device = torch.device('npu')
            else:
                device = torch.device('cpu')

            # 修复：强制移动到 NPU
            self.to(device)

        def forward(self, input):
            noise_input, conditional_input = input

            inputs = torch.cat((noise_input, conditional_input), dim=1)
            # 解决npu非连续内存问题
            inputs = inputs.contiguous()

            # 扩展维度
            expanded = self.expand(inputs)  # [B, 4]

            output = self.output(expanded)
            return output

    class Discriminator(nn.Module):
        """A simple discriminator for testing."""

        def __init__(self, data_input_shape, conditional_input_shape):
            super(Discriminator, self).__init__()
            self.data_input_shape = data_input_shape
            self.conditional_input_shape = conditional_input_shape

            data_dim = data_input_shape[
                1:]  # Extracting the actual data dimension
            conditional_dim = conditional_input_shape[
                1:]  # Extracting the actual conditional dimension
            input_dim = sum(data_dim) + sum(conditional_dim)

            # Define the dense layers
            self.dense1 = nn.Linear(input_dim, 10)
            self.dense2 = nn.Linear(10, 1)

        def forward(self, input):
            data_input, conditional_input = input
            discrim_in = torch.cat((data_input, conditional_input), dim=1)
            x = F.relu(self.dense1(discrim_in))
            a = self.dense2(x)
            output = torch.sigmoid(a)
            return output

    class ExampleGANModel(GANModel):
        """A simple GAN for testing."""

        def get_noise_input_shape(self):
            return (
                100,
                2,
            )

        def get_data_input_shapes(self):
            return [(
                100,
                1,
            )]

        def get_conditional_input_shapes(self):
            return [(
                100,
                1,
            )]

        def create_generator(self):
            noise_dim = self.get_noise_input_shape()
            conditional_dim = self.get_conditional_input_shapes()[0]

            return nn.Sequential(Generator(noise_dim, conditional_dim))

        def create_discriminator(self):
            data_input_shape = self.get_data_input_shapes()[0]
            conditional_input_shape = self.get_conditional_input_shapes()[0]

            return nn.Sequential(
                Discriminator(data_input_shape, conditional_input_shape))

        # 必须重写：支持 NPU
        def __init__(self,
                    n_generators: int = 1,
                    n_discriminators: int = 1,
                    create_discriminator_loss: Optional[Callable] = None,
                    create_generator_loss: Optional[Callable] = None,
                    _call_discriminator: Optional[Callable] = None,
                    device: Optional[Union[str, torch.device]] = None,
                    **kwargs):
        
            # 处理 device：支持 'npu'
            if device is None:
                if torch.cuda.is_available():
                    self.device = torch.device('cuda')
                elif torch.backends.mps.is_available():
                    self.device = torch.device('mps')
                elif is_npu_available():  # 关键：支持 NPU
                    self.device = torch.device('npu')
                else:
                    self.device = torch.device('cpu')
            else:
                if isinstance(device, str):
                    if device.lower() == 'npu':
                        if not torch.npu.is_available():
                            raise RuntimeError("NPU is not available. Please install CANN and torch.npu.")
                        self.device = torch.device('npu')
                    else:
                        self.device = torch.device(device)
                else:
                    self.device = device
            
            print(f"模型所在设备: {self.device}")

            # 保存参数
            self.n_generators = n_generators
            self.n_discriminators = n_discriminators

            # 调用父类构造函数（传入 device）
            super().__init__(
                n_generators=n_generators,
                n_discriminators=n_discriminators,
                create_discriminator_loss=create_discriminator_loss,
                create_generator_loss=create_generator_loss,
                _call_discriminator=_call_discriminator,
                device=self.device,
                **kwargs
            )

    has_torch = True
except ModuleNotFoundError:
    has_torch = False

def create_generator(noise_dim, conditional_dim):
    noise_dim = noise_dim
    conditional_dim = conditional_dim[0]
    return nn.Sequential(Generator(noise_dim, conditional_dim))

def create_discriminator(data_input_shape, conditional_input_shape):
    data_input_shape = data_input_shape[0]
    conditional_input_shape = conditional_input_shape[0]
    return nn.Sequential(
        Discriminator(data_input_shape, conditional_input_shape))

def generate_batch(batch_size):
    """Draw training data from a Gaussian distribution, where the mean  is a conditional input."""
    means = 10 * np.random.random([batch_size, 1])
    values = np.random.normal(means, scale=2.0)
    return means, values

def generate_data(gan, batches, batch_size):
    for _ in range(batches):
        means, values = generate_batch(batch_size)
        batch = {gan.data_inputs[0]: values, gan.conditional_inputs[0]: means}
        yield batch

gan = ExampleGANModel(learning_rate=0.01)
data = generate_data(gan, 500, 100)
gan.fit_gan(data, generator_steps=0.5, checkpoint_interval=0)

# See if it has done a plausible job of learning the distribution.

means = 10 * np.random.random([1000, 1])
values = gan.predict_gan_generator(conditional_inputs=[means])

测试结果：