RSAM2layerLOC2.py 文件源码

def train(epoch,lr,f):
    model.train()
    train_loss=0
    for batch_idx, (data, target) in enumerate(train_loader):
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr
        optimizer.zero_grad()
        output,l=model(data,target)
        loss = F.nll_loss(output[-1], target)
        if not args.lastout:
            for i in range(seq-1):
                loss = F.nll_loss(output[i], target)+loss
        loss.backward()
        optimizer.step()
        train_loss+=loss.data[0]
        if epoch>=0:
            torchvision.utils.save_image(l,'{}.png'.format(epoch))
        if batch_idx % args.log_interval == 0 and args.verbose:
            #print(l,file=f)
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.data[0]))
    train_loss /= len(train_loader) 
    return train_loss