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)