RIPGeo代码理解（六）main.py（运行模型进行训练和测试）

代码链接：RIPGeo代码实现

   ├── preprocess.py # 预处理数据集并为模型运行执行IP聚类
   ├── main.py # 运行模型进行训练和测试
   ├── test.py #加载检查点，然后测试

一、导入各种模块和数据库

import torch.nn

from lib.utils import *
import argparse
import numpy as np
import random, os
from lib.model import *
# import wandb
import copy

整体功能是准备运行一个 PyTorch 深度学习模型的环境，具体的功能实现需要查看 lib.utils、lib.model 中的代码，以及整个文件的后续部分。

1、from lib.utils import *：从 lib.utils 模块中导入所有内容。

2、from lib.model import *：从 lib.model 模块中导入所有内容。

3、import copy：导入 copy 模块，用于复制对象，通常用于创建对象的深拷贝。

二、参数初始化（通过命令行参数）

parser = argparse.ArgumentParser()
# parameters of initializing
parser.add_argument('--seed', type=int, default=1024, help='manual seed')
parser.add_argument('--model_name', type=str, default='RIPGeo')
parser.add_argument('--dataset', type=str, default='New_York', choices=["Shanghai", "New_York", "Los_Angeles"],
                    help='which dataset to use')

这部分代码的目的是通过命令行参数设置一些初始化的参数，例如随机数种子、模型名称和数据集名称。这使得在运行脚本时可以通过命令行参数来指定这些参数的值。

1、parser = argparse.ArgumentParser()：创建一个 argparse.ArgumentParser 对象，用于解析命令行参数。

2、parser.add_argument('--seed', type=int, default=1024, help='manual seed')：添加一个命令行参数，名称为 '--seed'，表示随机数种子，类型为整数，默认值为 1024，help 参数是在命令行中输入 --help 时显示的帮助信息。

3、parser.add_argument('--model_name', type=str, default='RIPGeo')：添加一个命令行参数，名称为 '--model_name'，表示模型的名称，类型为字符串，默认值为 'RIPGeo'。

4、parser.add_argument('--dataset', type=str, default='New_York', choices=["Shanghai", "New_York", "Los_Angeles"], help='which dataset to use')：添加一个命令行参数，名称为 '--dataset'，表示数据集的名称，类型为字符串，默认值为 'New_York'，choices 参数指定了可选的值为 ["Shanghai", "New_York", "Los_Angeles"]，用户只能从这三个值中选择。

三、训练过程参数设置

# parameters of training
parser.add_argument('--beta1', type=float, default=0.9)
parser.add_argument('--beta2', type=float, default=0.999)

parser.add_argument('--lr', type=float, default=2e-3)
parser.add_argument('--harved_epoch', type=int, default=5)
parser.add_argument('--early_stop_epoch', type=int, default=50)
parser.add_argument('--saved_epoch', type=int, default=100)

这部分代码的目的是设置一些训练过程中的超参数，例如优化器的动量参数、学习率、权重参数等。这些参数在训练过程中会影响模型的更新和收敛速度。

1、parser.add_argument('--beta1', type=float, default=0.9)：添加一个命令行参数，名称为 '--beta1'，表示 Adam 优化器的第一个动量（momentum）参数，类型为浮点数，默认值为 0.9。

2、parser.add_argument('--beta2', type=float, default=0.999)：添加一个命令行参数，名称为 '--beta2'，表示 Adam 优化器的第二个动量参数，类型为浮点数，默认值为 0.999。

3、parser.add_argument('--lr', type=float, default=2e-3)：添加一个命令行参数，名称为 '--lr'，表示学习率，类型为浮点数，默认值为 2e-3。

4、parser.add_argument('--harved_epoch', type=int, default=5)：添加一个命令行参数，名称为 '--harved_epoch'，表示当连续多少个epoch的性能没有增加时，学习率减半，类型为整数，默认值为 5。

5、parser.add_argument('--early_stop_epoch', type=int, default=50)：添加一个命令行参数，名称为 '--early_stop_epoch'，表示当连续多少个epoch的性能没有增加时，训练停止，类型为整数，默认值为 50。

6、parser.add_argument('--saved_epoch', type=int, default=100)：添加一个命令行参数，名称为 '--saved_epoch'，表示为测试保存多少个checkpoint（epoch），类型为整数，默认值为 100。

四、模型参数设置

# parameters of model
parser.add_argument('--dim_in', type=int, default=30, choices=[51, 30], help="51 if Shanghai / 30 else")
parser.add_argument('--dim_med', type=int, default=32)
parser.add_argument('--dim_z', type=int, default=32)
parser.add_argument('--eta', type=float, default=0.1)
parser.add_argument('--zeta', type=float, default=0.1)
parser.add_argument('--step', type=int, default=2)
parser.add_argument('--mu', type=float, default=0.2)
parser.add_argument('--lambda_1', type=float, default=1)
parser.add_argument('--lambda_2', type=float, default=1)
parser.add_argument('--c_mlp', type=bool, default=True)
parser.add_argument('--epoch_threshold', type=int, default=50)

opt = parser.parse_args()

这部分用于定义模型的结构和训练过程中的一些重要参数。

1、parser.add_argument('--dim_in', type=int, default=30, choices=[51, 30], help="51 if Shanghai / 30 else")：添加一个命令行参数，名称为 ''--dim_in'，表示输入数据的维度，类型为整数，默认值为 30，可选的有[51,30]，如果是上海数据集，维度为51，否则为30。

2、parser.add_argument('--dim_med', type=int, default=32)：添加一个命令行参数，名称为 '--dim_med'，表示中间层的维度，类型为整数，默认值为 32。

3、parser.add_argument('--dim_z', type=int, default=32)：添加一个命令行参数，名称为 '--dim_z'，表示向量表示的维度，类型为整数，默认值为 32。

4、parser.add_argument('--eta', type=float, default=0.1)：添加一个命令行参数，名称为 '--eta'，表示数据扰动程度，默认值为 0.1。

5、parser.add_argument('--zeta', type=float, default=0.1)：添加一个命令行参数，名称为 '--zeta'，表示参数扰动程度，默认值为 0.1。

6、parser.add_argument('--step', type=int, default=2)：添加一个命令行参数，名称为 '--step'，表示单参数扰动下梯度上升次数，类型为整数，默认值为 2。

7、parser.add_argument('--mu', type=float, default=0.2)：添加一个命令行参数，名称为 '--mu'，表示参数扰动的内学习率，默认值为 0.2。

8、parser.add_argument('--lambda_1', type=float, default=1)：添加一个命令行参数，名称为 '--lambda_1'，表示损失函数中数据扰动的权衡系数，默认值为 1。

9、parser.add_argument('--lambda_2', type=float, default=1)：添加一个命令行参数，名称为 '--lambda_2'，表示损失函数中参数扰动的权衡系数，默认值为 1。

10、parser.add_argument('--c_mlp', type=bool, default=True)：添加一个命令行参数，名称为 '--c_mlp'，表示在预测是否使用collaborative_mlp时，默认值为 True。

11、parser.add_argument('--epoch_threshold', type=int, default=50)：添加一个命令行参数，名称为 '--epoch_threshold'，表示当我们开始在数据和参数中添加扰动时，类型为整数，默认值为 50。

12、opt = parser.parse_args()：将命令行参数解析成Python对象。简单来说，就是通过parser解析命令行传入的参数，并将其赋值给变量pt。

五、设置随机种子数

if opt.seed:
    print("Random Seed: ", opt.seed)
    random.seed(opt.seed)
    torch.manual_seed(opt.seed)
torch.set_printoptions(threshold=float('inf'))

这一部分的目的是确保在使用随机数的场景中，每次运行程序得到的随机结果是可复现的。通过设置相同的随机数种子，可以使得每次运行得到相同的随机数序列。

1、如果 opt 对象中的 seed 属性存在（不为 0 或 False 等假值），则执行以下操作：

打印随机数种子的信息。
使用 random 模块设置 Python 内建的随机数生成器的种子。
使用 PyTorch 的 torch 模块设置随机数种子。

2、torch.set_printoptions(threshold=float('inf'))：设置 PyTorch 的打印选项，将打印的元素数量限制设置为无穷大，即不限制打印的元素数量。这样可以确保在打印张量时，所有元素都会被打印出来，而不会被省略。

六、过滤所有警告信息

warnings.filterwarnings('ignore')

过滤掉所有警告信息，将警告信息忽略。这通常用于在代码中避免显示一些不影响程序执行的警告信息，以保持输出的清晰。在某些情况下，警告信息可能是有用的，但如果明确知道这些警告对程序执行没有影响，可以选择忽略它们。

七、动态选择运行环境

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print("device:", device)
print("Dataset: ", opt.dataset)
cuda = True if torch.cuda.is_available() else False
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor

这部分代码的目的是根据硬件环境动态选择运行模型的设备，并选择相应的 PyTorch 张量类型。如果有可用的 GPU，就使用 GPU 运行模型和 GPU 张量类型；否则，使用 CPU 运行模型和 CPU 张量类型。

1、device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')：创建一个 PyTorch 设备对象，表示运行模型的设备。如果 CUDA 可用（即有可用的 GPU），则使用 'cuda:0' 表示第一个 GPU，否则使用 'cpu' 表示 CPU。

2、print("device:", device)：打印设备的信息，即使用的是 GPU 还是 CPU。

3、cuda = True if torch.cuda.is_available() else False：根据 CUDA 是否可用设置一个布尔值，表示是否使用 GPU。如果 CUDA 可用，则 cuda 为 True，否则为 False。

4、Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor：根据上一步得到的 cuda 布尔值选择使用 GPU 还是 CPU 上的 PyTorch 张量类型。如果 cuda 为 True，则 Tensor 被设置为 torch.cuda.FloatTensor，表示在 GPU 上的浮点数张量类型，否则设置为 torch.FloatTensor，表示在 CPU 上的浮点数张量类型。

八、加载数据（训练测试）

'''load data'''
train_data = np.load("./datasets/{}/Clustering_s1234_lm70_train.npz".format(opt.dataset),
                     allow_pickle=True)
test_data = np.load("./datasets/{}/Clustering_s1234_lm70_test.npz".format(opt.dataset),
                    allow_pickle=True)
train_data, test_data = train_data["data"], test_data["data"]
print("data loaded.")

这部分代码的目的是加载训练集和测试集的数据，数据文件的路径根据 opt.dataset 的值确定（见四、模型参数设置）。

train_data = np.load("./datasets/{}/Clustering_s1234_lm70_train.npz".format(opt.dataset), allow_pickle=True)：使用 NumPy 的 load 函数加载训练数据。数据集的路径根据opt.dataset的取值而动态确定。allo