ita9naiwa

'''
Train CIFAR10 with PyTorch.
based on https://github.com/kuangliu/pytorch-cifar
'''
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn

ita9naiwa

ita9naiwa

class GATMHA(nn.Module):
    def __init__(self, hidden_size, n_heads, lralpha=0.2):
        super(GATMHA, self).__init__()
        self.W = nn.Linear(hidden_size, hidden_size, bias=True)
        self.Q = nn.Linear(hidden_size, hidden_size, bias=True)
        self.a = nn.Parameter(torch.FloatTensor(n_heads, 2 * (hidden_size // n_heads)))
        self.lralpha = lralpha
        self.n_heads = n_heads
        self.n_hidden_per_head = hidden_size // n_heads
        nn.init.uniform_(self.W.weight, -1 / np.sqrt(hidden_size), 1 / np.sqrt(hidden_size))

ita9naiwa

ita9naiwa

ita9naiwa

# mainly brought from https://github.com/pytorch/examples/blob/master/reinforcement_learning/reinforce.py
import gym
import numpy as np
import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torch.distributions import Categorical
import torch.multiprocessing as mp