加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
该仓库未声明开源许可证文件(LICENSE),使用请关注具体项目描述及其代码上游依赖。
克隆/下载
utils_model.py 1.07 KB
一键复制 编辑 原始数据 按行查看 历史
zeng798473532 提交于 2021-11-07 13:59 . initial
import torch
import torch.nn as nn
import torch.nn.functional as F
class DQN(nn.Module):
def __init__(self, action_dim, device):
super(DQN, self).__init__()
self.__conv1 = nn.Conv2d(4, 32, kernel_size=8, stride=4, bias=False)
self.__conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2, bias=False)
self.__conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1, bias=False)
self.__fc1 = nn.Linear(64*7*7, 512)
self.__fc2 = nn.Linear(512, action_dim)
self.__device = device
def forward(self, x):
x = x / 255.
x = F.relu(self.__conv1(x))
x = F.relu(self.__conv2(x))
x = F.relu(self.__conv3(x))
x = F.relu(self.__fc1(x.view(x.size(0), -1)))
return self.__fc2(x)
@staticmethod
def init_weights(module):
if isinstance(module, nn.Linear):
torch.nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
module.bias.data.fill_(0.0)
elif isinstance(module, nn.Conv2d):
torch.nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化