avg_D_real_m_loss = 0 avg_D_real_m2_loss = 0 avg_D_fake_loss = 0 avg_G_fake_loss = 0 avg_percept_loss = 0 for i ,(img,att,seg,cat,nnseg) in enumerate(train_loader): bs = img.size(0) rnd_batch_num = np.random.randint(len(train_data),size=bs) rnd_att_list = [train_data[i][1] for i in rnd_batch_num] rnd_att_np = np.asarray(rnd_att_list) rnd_att = torch.from_numpy(rnd_att_np).float() #convert images to tensors and send to gpu seg = seg.type(torch.FloatTensor) nnseg = nnseg.type(torch.FloatTensor) img = Variable(img.cuda()) att = Variable(att.cuda()) rnd_att = Variable(rnd_att.cuda()) seg = Variable(seg.cuda()) nnseg = Variable(nnseg.cuda()) cat = Variable(cat.cuda()) Z = init_z_foreach_layout(cat, bs) img_norm = img * 2 - 1 img_G = img_norm requires_grad(G, False) requires_grad(D, True) D.zero_grad() #calculate loss for real image with segmask and attributes
def train_model(taskid_path, args, net, train_dataloader, val_dataloader, density, dataset_length): log_path = os.path.join(taskid_path, "log") make_dir(log_path) writer = SummaryWriter(log_path) torch.cuda.set_device(args.gpu_id) net = net.cuda() max_acc = 0 max_epoch = 0 judge = 0 for epoch in range(args.epochnum): paras = dict(net.named_parameters()) paras_new = [] for k, v in paras.items(): if 'mask' in k: if 'bias' in k: paras_new += [{ 'params': [v], 'lr': args.lr[epoch] * 2, 'weight_decay': args.weightdecay * 0 }] if 'mask_weight' in k: paras_new += [{ 'params': [v], 'lr': args.lr[epoch] * 0.05, 'weight_decay': args.weightdecay * 0 }] if '.weight' in k: paras_new += [{ 'params': [v], 'lr': args.lr[epoch] * 1, 'weight_decay': args.weightdecay * 1 }] if 'line' in k: if 'bias' in k: paras_new += [{ 'params': [v], 'lr': args.lr[epoch] * 2, 'weight_decay': args.weightdecay * 0 }] if 'weight' in k: paras_new += [{ 'params': [v], 'lr': args.lr[epoch] * 1, 'weight_decay': args.weightdecay * 1 }] if 'conv' in k: if 'bias' in k: paras_new += [{ 'params': [v], 'lr': args.lr[epoch] * 1, 'weight_decay': args.weightdecay * 1 }] if 'weight' in k: paras_new += [{ 'params': [v], 'lr': args.lr[epoch] * 1, 'weight_decay': args.weightdecay * 1 }] optimizer = SGD(paras_new, lr=args.lr[epoch], momentum=args.momentum, weight_decay=args.weightdecay) # train net.train() train_loss = [] train_acc = [] print('Train: ' + "\n" + 'epoch:{}'.format(epoch + 1)) for index, (image, label) in enumerate(train_dataloader): batch_size = image.shape[0] image = Variable(image) image = image.cuda() label = label.cuda() out = net(image, label, torch.Tensor([epoch + 1]), density) if args.model == "resnet_18" or args.model == "resnet_50" or args.model == "densenet_121": out = torch.unsqueeze(out, 2) out = torch.unsqueeze(out, 3) label = Variable(label) if args.losstype == 'logistic': loss = logistic_F.apply(out, label) train_loss.append(loss.cpu().clone().data.numpy()) train_correct = label.mul(out) train_correct = torch.max( train_correct, torch.zeros(train_correct.size()).cuda()) train_correct = torch.sum((train_correct > 0)) train_acc.append(train_correct.cpu().data.numpy()) if args.losstype == 'softmax': loss = softmax_F.apply(out, label) train_loss.append(loss.cpu().clone().data.numpy()) (tmp, out) = torch.sort(out, dim=1, descending=True) (tmp, label) = torch.max(label, dim=1) label = label.unsqueeze(2) error = ~(out == label) train_correct = args.batchsize - torch.sum(error[:, 0, 0, 0]) train_acc.append(train_correct.cpu().data.numpy()) optimizer.zero_grad() loss.backward() optimizer.step() print('batch:{}/{}'.format(index + 1, len(train_dataloader)) + " " + 'loss:{:.6f}'.format(loss / batch_size) + " " + 'acc:{:.6f}'.format(train_correct.cpu().data.numpy() / (batch_size * args.label_num))) length = dataset_length['train'] if index + 1 == len( train_dataloader) else args.batchsize * (index + 1) if (index + 1) % 10: writer.add_scalar('Train/Loss', sum(train_loss) / length, epoch) writer.add_scalar('Train/acc', sum(train_acc) / (length * args.label_num), epoch) # eval net.eval() with torch.no_grad(): eval_loss = [] eval_acc = [] for index, (image, label) in enumerate(val_dataloader): print('Val: ' + "\n" + 'epoch:{}'.format(epoch + 1)) batch_size = image.shape[0] image = Variable(image) image = image.cuda() label = label.cuda() out = net(image, label, torch.Tensor([epoch + 1]), density) if args.model == "resnet_18" or args.model == "resnet_50" or args.model == "densenet_121": out = torch.unsqueeze(out, 2) out = torch.unsqueeze(out, 3) label = Variable(label) if args.losstype == 'logistic': loss = logistic_F.apply(out, label) eval_loss.append(loss.cpu().data.numpy()) eval_correct = label.mul(out) eval_correct = torch.max( eval_correct, torch.zeros(eval_correct.size()).cuda()) eval_correct = torch.sum((eval_correct > 0)) eval_acc.append(eval_correct.cpu().data.numpy()) if args.losstype == 'softmax': loss = softmax_F.apply(out, label) eval_loss.append(loss.cpu().data.numpy()) (tmp, out) = torch.sort(out, dim=1, descending=True) (tmp, label) = torch.max(label, dim=1) label = label.unsqueeze(2) error = ~(out == label) eval_correct = args.batchsize - torch.sum(error[:, 0, 0, 0]) eval_acc.append(eval_correct.cpu().data.numpy()) length = dataset_length['val'] if index + 1 == len( val_dataloader) else args.batchsize * (index + 1) print('batch:{}/{}'.format(index + 1, len(val_dataloader)) + " " + 'loss:{:.6f}'.format(loss / batch_size) + " " + 'acc:{:.6f}'.format(eval_correct.cpu().data.numpy() / (batch_size * args.label_num))) print("max_acc:" + str(max_acc)) if sum(eval_acc) / (length * args.label_num) > max_acc: judge = 1 max_acc = sum(eval_acc) / (length * args.label_num) print("rightnow max_acc:" + str(max_acc)) max_epoch = epoch writer.add_scalar('Eval/Loss', sum(eval_loss) / length, epoch) writer.add_scalar('Eval/acc', sum(eval_acc) / (length * args.label_num), epoch) if judge == 1 or (epoch + 1) % 50 == 0: # save torch.save(net, taskid_path + '/net-' + str(epoch + 1) + '.pkl') #torch.save(net.state_dict(), taskid_path + '/net-params-' + str(epoch + 1) + '.pkl') judge = 0 return max_acc, max_epoch + 1