def load_single_model(self, trained_model_filepath):
'''
func: 加载并实例化trained_model_filepath这个路径下的模型
Parameter
---------
trained_model_filepath: str
训练好的单帧预测模型的保存路径，
eg: '/home/All_12_Basic_model/Basic0/checkpoint.chk'
Returns:
--------
network: model
已经加载并实例化并后的模型
'''
from leibniz.unet.base import UNet
from leibniz.unet.residual import Basic
from leibniz.nn.activation import CappingRelu
network = UNet(in_channels=self.in_dim,
out_channels=self.single_out_dim,
normalizor='batch',
spatial=(256, 256),
layers=4,
ratio=0,
vblks=[2, 2, 2, 2],
hblks=[0, 0, 0, 0],
scales=[-1, -1, -1, -1],
factors=[1, 1, 1, 1],
block=Basic,
relu=CappingRelu(),
final_normalized=False)
checkpoint = torch.load(trained_model_filepath)
network.load_state_dict(checkpoint['net'])
return network.to(self.device)
def __init__(self):
super().__init__()
self.relu = nn.ReLU()
self.enc = resunet(10,
62,
block=HyperBottleneck,
relu=nn.ReLU(inplace=True),
layers=6,
ratio=-2,
vblks=[1, 1, 1, 1, 1, 1],
hblks=[3, 3, 3, 3, 3, 3],
scales=[-1, -1, -1, -1, -1, -1],
factors=[1, 1, 1, 1, 1, 1],
spatial=(64, 64))
self.dec = resunet(10,
10,
block=HyperBottleneck,
relu=CappingRelu(),
layers=6,
ratio=-3,
vblks=[1, 1, 1, 1, 1, 1],
hblks=[0, 0, 0, 0, 0, 0],
scales=[-1, -1, -1, -1, -1, -1],
factors=[1, 1, 1, 1, 1, 1],
spatial=(64, 64),
final_normalized=True)
def __init__(self, setting):
super().__init__(setting)
self.name = 't850d3bg-rasp'
self.resunet = resunet(setting.input_span * len(setting.vars_in) + self.constant_size + 2, 1,
spatial=(32, 64+2), layers=5, ratio=-2,
vblks=[2, 2, 2, 2, 2], hblks=[1, 1, 1, 1, 1],
scales=[-1, -1, -1, -1, -1], factors=[1, 1, 1, 1, 1],
block=HyperBottleneck, relu=CappingRelu(), final_normalized=False)
def __init__(self):
super().__init__()
self.unet = UNet(2,
10,
normalizor='batch',
spatial=(32, 32),
layers=5,
ratio=0,
vblks=[2, 2, 2, 2, 2],
hblks=[2, 2, 2, 2, 2],
scales=[-1, -1, -1, -1, -1],
factors=[1, 1, 1, 1, 1],
block=DirectBlocks,
relu=CappingRelu(),
final_normalized=True)
def __init__(self):
super().__init__()
self.tube = hyptub(10,
25,
10,
encoder=resunet,
decoder=resunet,
block=HyperBottleneck,
relu=CappingRelu(),
ratio=-1.5,
layers=6,
vblks=[1, 1, 1, 1, 1, 1],
hblks=[1, 1, 1, 1, 1, 1],
scales=[-1, -1, -1, -1, -1, -1],
factors=[1, 1, 1, 1, 1, 1],
spatial=(64, 64))
def __init__(self, setting):
super().__init__(setting)
tube = hyptub(1664, 832, 1664, encoder=linear, decoder=linear)
self.resunet = resunet(setting.input_span * (len(setting.vars) + 2) +
self.constant_size + 2,
1,
spatial=(32, 64 + 2),
layers=5,
ratio=-1,
vblks=[2, 2, 2, 2, 2],
hblks=[1, 1, 1, 1, 1],
scales=[-1, -1, -1, -1, -1],
factors=[1, 1, 1, 1, 1],
block=HyperBottleneck,
relu=CappingRelu(),
enhencer=tube,
final_normalized=False)
def __init__(self):
super().__init__()
if configs.add_constants:
self.constants = get_constants()
self.unet = UNet(configs.input_dim,
configs.output_dim,
block=HyperBottleneck,
relu=CappingRelu(),
layers=4,
ratio=0,
vblks=[2, 2, 2, 2],
hblks=[0, 0, 0, 0],
scales=[-1, -1, -1, -1],
factors=[1, 1, 1, 1],
spatial=(32, 64),
normalizor='batch',
padding=None,
final_normalized=False)
def __init__(self):
super().__init__()
if configs.add_constants:
self.constants = get_constants()
self.unet = ReUNet(configs.input_dim,
configs.output_dim,
block=WarpBottleneck,
relu=CappingRelu(),
layers=4,
ratio=0,
vblks=[2, 2, 2, 2],
hblks=[0, 0, 0, 0],
scales=[s, s, s, s],
factors=[1, 1, 1, 1],
spatial=(6, 48, 48),
normalizor='batch',
padding=CubeSpherePadding2D(1),
final_normalized=False)
else:
model_configs['out_channels'] = dataset_configs['output_frames']
model_configs['normalizor'] = 'batch'
model_configs['spatial'] = dataset_configs['aim_size']
model_configs['layers'] = 4
model_configs['ratio'] = 0
model_configs['vblks'] = [2] * model_configs['layers']
model_configs['hblks'] = [0] * model_configs['layers']
model_configs['scales'] = [-1] * model_configs['layers']
model_configs['factors'] = [1] * model_configs['layers']
# model_configs['block'] = wp.WarpBottleneck
model_configs['block'] = Basic
# model_configs['block'] = HyperBottleneck
model_configs['relu'] = CappingRelu()
model_configs['final_normalized'] = False
network = UNet(**model_configs)
#%%
##损失函数及其参数配置
loss_configs = {}
loss_configs['a'] = 6
loss_configs['b'] = 0.8
loss_configs['mse_w'] = 1
loss_configs['mae_w'] = 1
loss_fn = BLoss.BEXPMSAELoss(**loss_configs)
loss_configs['loss_name'] = loss_fn.__str__()[:-2]
