Why does Cage visit the camp after the Omega is dead?

At the end Cage was not in the camp as a trainee or deserter, but as a major. When he awoke in the helicopter after killing the Omega, the time was still before his encounter with General Brigham. So at this time he still had his complete military status as Major Cage, since there was no need to go into the field anymore and thus no need to desert at all.

The only reason why he was at the camp was to seek out Rita Vrataski, but not because he had to prepare her for the upcoming fight as he did hundreds of days before, but because against her advice he grew a strong affection, if not love, for her. He simply needed to tell her what an amazing adventure they just experienced, though knowing that she certainly didn't experience it at all, which is also the reason why he smiles, trying to find the words to convice her of all they've been through, just this time under a happier premise. And in fact he had just seen her die, so first of all, he needed to make sure she was still alive, even if knowing it.

He's in the camp initially because he was arrested and forced into front line duty against his will (if you can even say that for a military position), but the reason for why he went back at the very end of the movie, after humanity won, can be explained in different terms:

1. It's what he'd been doing all along. I don't recall if Cage ever mentioned how many times he'd died and "looped back", but every single time he did he went to Rita for training, in addition to receiving her help in every attempt at finding the Omega. You could sort of look at it as, "Old habits die hard."

2. To provide Rita with closure. Rita knows that it's possible for a human to absorb the time reset powers of an Omega through an Alpha, so Cage has to return to her so she'll know it was him that caused their victory, and to explain why. Seriously, she had the ability herself and was attempting to loop through the same day to find the Omega and kill it before she lost the ability, because she knew that's what needed to be done in order win, and suddenly humanity wins against the Mimics. I know I'd want to know how the hell it happened if I was the only one (besides Carter) who knew that was possible, and thus why humanity was ultimately doomed to lose unless someone else achieved what she had.