models.py 文件源码

def embed_state(self, previous_states, state, volatile=False, requires_grad=True, gpu=-1):
        prev_scrs = [self.downscale_prev(s.screenshot_rs) for s in previous_states]
        prev_scrs_y = [cv2.cvtColor(scr, cv2.COLOR_RGB2GRAY) for scr in prev_scrs]

        #inputs = np.dstack([self.downscale(state.screenshot_rs)] + list(reversed(prev_scrs_y)))
        inputs = np.array(self.downscale(state.screenshot_rs), dtype=np.float32)
        inputs = inputs / 255.0
        inputs = inputs.transpose((2, 0, 1))
        inputs = inputs[np.newaxis, ...]
        inputs = to_cuda(to_variable(inputs, volatile=volatile, requires_grad=requires_grad), gpu)

        inputs_prev = np.dstack(prev_scrs_y)
        inputs_prev = inputs_prev.astype(np.float32) / 255.0
        inputs_prev = inputs_prev.transpose((2, 0, 1))
        inputs_prev = inputs_prev[np.newaxis, ...]
        inputs_prev = to_cuda(to_variable(inputs_prev, volatile=volatile, requires_grad=requires_grad), gpu)

        return self.embed(inputs, inputs_prev)