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

docker exec -it {container_name} bash

pip install deepchem-ascend==0.0.1

pip install torch-geometric

import numpy as np
import deepchem as dc
import tempfile
import os

try:
    import torch
    import torch.nn as nn
    from deepchem.models.torch_models import ProgressiveMultitaskRegressor, ProgressiveMultitaskClassifier
    has_torch = True
except ModuleNotFoundError:
    has_torch = False
    pass

n_tasks = 2
n_features = 12

torch_model = ProgressiveMultitaskRegressor(
    n_tasks=n_tasks,
    n_features=n_features,
    layer_sizes=[128, 256],
    alpha_init_stddevs=0.02,
    weight_init_stddevs=0.02,
    dropouts=0,
)

weights = np.load(
    os.path.join(os.path.dirname(__file__), "assets",
                    "progressive-multitask-regressor-sample-weights.npz"))

# Determine if torch_model is a DeepChem TorchModel or raw nn.Module
if hasattr(torch_model, 'model'):
    # DeepChem's TorchModel wraps the actual nn.Module in .model
    nn_model = torch_model.model
else:
    # Assume it's already an nn.Module
    nn_model = torch_model

# Get the device of the model (e.g., 'npu:0' or 'cpu')
device = next(nn_model.parameters()).device

def to_tensor_on_device(np_array):
    """Convert numpy array to torch tensor on the model's device."""
    return torch.from_numpy(np_array).to(device)

# === Load layer weights for task 0 ===
nn_model.layers[0][0].weight.data.copy_(to_tensor_on_device(weights["layer-0-0-w"]))
nn_model.layers[0][0].bias.data.copy_(to_tensor_on_device(weights["layer-0-0-b"]))

nn_model.layers[0][1].weight.data.copy_(to_tensor_on_device(weights["layer-0-1-w"]))
nn_model.layers[0][1].bias.data.copy_(to_tensor_on_device(weights["layer-0-1-b"]))

nn_model.layers[0][2].weight.data.copy_(to_tensor_on_device(weights["layer-0-2-w"]))
nn_model.layers[0][2].bias.data.copy_(to_tensor_on_device(weights["layer-0-2-b"]))

# === Load layer weights for task 1 ===
nn_model.layers[1][0].weight.data.copy_(to_tensor_on_device(weights["layer-1-0-w"]))
nn_model.layers[1][0].bias.data.copy_(to_tensor_on_device(weights["layer-1-0-b"]))

nn_model.layers[1][1].weight.data.copy_(to_tensor_on_device(weights["layer-1-1-w"]))
nn_model.layers[1][1].bias.data.copy_(to_tensor_on_device(weights["layer-1-1-b"]))

nn_model.layers[1][2].weight.data.copy_(to_tensor_on_device(weights["layer-1-2-w"]))
nn_model.layers[1][2].bias.data.copy_(to_tensor_on_device(weights["layer-1-2-b"]))

# === Load adapter weights and alphas ===
# Task 1 adapter 0
nn_model.alphas[0][0].data.copy_(to_tensor_on_device(weights["alpha-0-0"]))
nn_model.adapters[0][0][0].weight.data.copy_(to_tensor_on_device(weights["adapter-0-0-0-w"]))
nn_model.adapters[0][0][0].bias.data.copy_(to_tensor_on_device(weights["adapter-0-0-0-b"]))
nn_model.adapters[0][0][1].weight.data.copy_(to_tensor_on_device(weights["adapter-0-0-1-w"]))

# Task 1 adapter 1
nn_model.alphas[0][1].data.copy_(to_tensor_on_device(weights["alpha-0-1"]))
nn_model.adapters[0][1][0].weight.data.copy_(to_tensor_on_device(weights["adapter-0-1-0-w"]))
nn_model.adapters[0][1][0].bias.data.copy_(to_tensor_on_device(weights["adapter-0-1-0-b"]))
nn_model.adapters[0][1][1].weight.data.copy_(to_tensor_on_device(weights["adapter-0-1-1-w"]))

input_x = weights["input"]
output = weights["output"]

# Inference using TorchModel's predict() method works with NumpyDataset only. Hence we need to convert our numpy arrays to NumpyDataset.
y = np.random.rand(input_x.shape[0], 1)
w = np.ones((input_x.shape[0], 1))
ids = np.arange(input_x.shape[0])
input_x = dc.data.NumpyDataset(input_x, y, w, ids)

torch_out = torch_model.predict(input_x)

print(f"模型所在设备: {device}")