def load_pretrain_weights(self):
"""Loading weights from trained MLP model & GMF model"""
config = self.config
config['latent_dim'] = config['latent_dim_mlp']
mlp_model = MLP(config)
if config['use_cuda'] is True:
mlp_model.cuda()
resume_checkpoint(mlp_model,
model_dir=config['pretrain_mlp'],
device_id=config['device_id'])
self.embedding_user_mlp.weight.data = mlp_model.embedding_user.weight.data
self.embedding_item_mlp.weight.data = mlp_model.embedding_item.weight.data
for idx in range(len(self.fc_layers)):
self.fc_layers[idx].weight.data = mlp_model.fc_layers[
idx].weight.data
config['latent_dim'] = config['latent_dim_mf']
gmf_model = GMF(config)
if config['use_cuda'] is True:
gmf_model.cuda()
resume_checkpoint(gmf_model,
model_dir=config['pretrain_mf'],
device_id=config['device_id'])
self.embedding_user_mf.weight.data = gmf_model.embedding_user.weight.data
self.embedding_item_mf.weight.data = gmf_model.embedding_item.weight.data
self.affine_output.weight.data = 0.5 * torch.cat([
mlp_model.affine_output.weight.data,
gmf_model.affine_output.weight.data
],
dim=-1)
self.affine_output.bias.data = 0.5 * (
mlp_model.affine_output.bias.data +
gmf_model.affine_output.bias.data)
def main(args):
#建立文件夹
result_path = os.path.join(args.output_path, "final_result")
check_dir(args.output_path)
check_dir(result_path)
#设置gpu
torch.cuda.set_device(args.gpu_id) #GPU
#加载数据
train_set = MySet(args.txt_path, mode="train")
val_set = MySet(args.txt_path, mode="val")
train_loader = DataLoader(train_set,
batch_size=args.batch_size,
shuffle=True)
val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False)
#网络模型
net = MLP(args)
net.cuda() #GPU
#使用Adam优化器和交叉熵损失函数
optimizer = torch.optim.Adam(net.parameters(),
lr=args.lr,
weight_decay=1e-4)
cost = torch.nn.CrossEntropyLoss()
epoch_loss, epoch_acc, val_loss, val_acc = [], [], [], []
#epoch循环
t0 = time.time()
for epoch in range(args.epoch):
t_loss, t_acc = train_epoch(net, train_loader, optimizer, cost)
epoch_loss.append(t_loss)
epoch_acc.append(t_acc)
v_loss, v_acc = val_epoch(net, val_loader, cost)
val_acc.append(v_acc)
val_loss.append(v_loss)
#画出训练集和测试集的loss及acc
plot(epoch_acc, val_acc, result_path, 'train_acc', 'val_acc',
args.batch_size, 'acc')
plot(epoch_loss, val_loss, result_path, 'train_loss', 'val_loss',
args.batch_size, 'loss')
info = [str(epoch).zfill(3), t_loss, v_acc]
print("Epoch: {} | train Loss: {:.4f} val ACC: {:.4f}".format(*info))
t1 = time.time()
print("Optimization Finished! Cost time:{:.1f} minutes".format(
(t1 - t0) / 60))
print("The final acc=%g" % v_acc)
#保存最终模型
state = net.state_dict()
torch.save(state, os.path.join(result_path, "Network_final.pth.gz"))
nusc_classes = [
'__background__', 'pedestrian', 'barrier', 'trafficcone', 'bicycle', 'bus',
'car', 'construction', 'motorcycle', 'trailer', 'truck'
]
batch_size = 4
nusc_sampler_batch = sampler(400, 2)
nusc_set = nuscenes_dataloader(batch_size, len(nusc_classes), training=True)
nusc_dataloader = torch.utils.data.DataLoader(nusc_set,
batch_size=batch_size,
sampler=nusc_sampler_batch)
nusc_iters_per_epoch = int(len(nusc_set) / batch_size)
num_epochs = 200
model = MLP()
model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
regressor = nn.SmoothL1Loss()
classifier = nn.BCELoss()
im = torch.FloatTensor(1)
points = torch.FloatTensor(1)
offset = torch.FloatTensor(1)
m = torch.FloatTensor(1)
rot_matrix = torch.FloatTensor(1)
gt_corners = torch.FloatTensor(1)
im = im.cuda()
points = points.cuda()
offset = offset.cuda()
import numpy as np
from sklearn.metrics import accuracy_score
gpu_id = 0
batch_size = 1000
parser = InitParser()
torch.cuda.set_device(gpu_id)
weights_path = "../output/final_result/Network_final.pth.gz"
test_set = MySet(parser.txt_path, mode="test")
net = MLP(parser)
weights = torch.load(weights_path, map_location='cuda:%d' % (gpu_id)) #GPU ###
net.load_state_dict(weights)
net.cuda().eval() #测试模式 GPU
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=True)
start = time.time()
labels, predicts = [], []
with torch.no_grad():
for batch_idx, (feature, label) in enumerate(test_loader):
feature = feature.float().cuda() #GPU
label = label.data.numpy()
predict = net(feature)
predict = predict.data.cpu().numpy()
predict = np.argmax(predict, axis=1)
labels.extend(list(label))
predicts.extend(list(predict))
acc = accuracy_score(labels, predicts)
